OpenScope’s Predictive Processing Mesoscope Dataset

import warnings
warnings.filterwarnings('ignore')

try:
    from databook_utils.dandi_utils import dandi_download_open, dandi_stream_open
except ImportError:
    !git clone https://github.com/AllenInstitute/openscope_databook.git
    %cd openscope_databook
    %pip install -e .

import math
import numpy as np
import pandas as pd
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from scipy import interpolate

%matplotlib inline

dandiset_id = "001768"
dandi_filepath = "sub-839909/sub-839909_ses-multiplane-ophys-839909-2026-02-20-12-53-27_ophys.nwb"
download_loc = "."

# This can take several minutes the first time; the file is cached locally afterward.
io = dandi_download_open(dandiset_id, dandi_filepath, download_loc) # use dandi_stream_open to stream the file without downloading it
nwb = io.read()

print(f"Session ID:          {nwb.session_id}")
print(f"Session description: {nwb.session_description}")
print(f"Session start time:  {nwb.session_start_time}")
print(f"Institution:         {nwb.institution}")
print(f"Experimenter:        {nwb.experimenter}")
print()

subj = nwb.subject
print(f"Subject ID:   {subj.subject_id}")
print(f"Species:      {subj.species}")
print(f"Sex:          {subj.sex}")
print(f"Age:          {subj.age}")
print(f"Genotype:     {subj.genotype}")

plane_names = sorted(nwb.imaging_planes.keys())
print(f"Number of imaging planes: {len(plane_names)}")
print()

for name in plane_names:
    plane = nwb.imaging_planes[name]
    print(f"  {name}:")
    print(f"    location:          {plane.location}")
    print(f"    imaging_rate:      {plane.imaging_rate} Hz")
    print(f"    excitation_lambda: {plane.excitation_lambda} nm")
    print(f"    indicator:         {plane.indicator}")
    if plane.grid_spacing is not None:
        print(f"    grid_spacing:      {plane.grid_spacing} {plane.grid_spacing_unit}")
    print()

plane_roi_stats = []

for name in plane_names:
    proc = nwb.processing[name]
    roi_table = proc["image_segmentation"]["roi_table"]
    n_total = len(roi_table)
    is_soma = np.array(roi_table["is_soma"].data)
    is_dendrite = np.array(roi_table["is_dendrite"].data)
    plane_roi_stats.append({
        "plane": name,
        "total_rois": n_total,
        "soma_rois": int(is_soma.sum()),
        "dendrite_rois": int(is_dendrite.sum()),
    })

roi_df = pd.DataFrame(plane_roi_stats)
roi_df["total_soma"] = roi_df["soma_rois"]
roi_df = roi_df.drop(columns="total_soma")

print(roi_df.to_string(index=False))
print()
print(f"Total ROIs across all planes: {roi_df['total_rois'].sum()}")
print(f"Total soma ROIs:              {roi_df['soma_rois'].sum()}")
print(f"Total dendrite ROIs:          {roi_df['dendrite_rois'].sum()}")

N_COLS_PLANES = 2   # how many planes side by side in the grid
n_planes = len(plane_names)
n_rows = math.ceil(n_planes / N_COLS_PLANES)

# Each plane occupies 2 axes: [max projection | ROI mask overlay]
fig, axes = plt.subplots(
    n_rows, N_COLS_PLANES * 2,
    figsize=(6 * N_COLS_PLANES * 2, 5 * n_rows)
)
axes = np.array(axes).reshape(n_rows, N_COLS_PLANES * 2)

for i, name in enumerate(plane_names):
    row      = i // N_COLS_PLANES
    col_base = (i %  N_COLS_PLANES) * 2   # left of the pair
    ax_max   = axes[row, col_base]
    ax_roi   = axes[row, col_base + 1]

    proc      = nwb.processing[name]
    image_mod = proc["images"]

    # Max projection (fall back to average if not present)
    if "max_projection" in image_mod.images:
        proj = np.array(image_mod["max_projection"].data)
        proj_label = "Max Projection"
    else:
        proj = np.array(image_mod["average_projection"].data)
        proj_label = "Avg Projection"

    roi_table = proc["image_segmentation"]["roi_table"]
    all_masks = np.array(roi_table["image_mask"].data[:])
    # Binarize: any non-zero weight pixel counts as part of an ROI
    binary_mask = (all_masks.max(axis=0) > 0).astype(np.float32)
    n_soma = int(np.array(roi_table["is_soma"].data).sum())

    ax_max.imshow(proj, cmap="gray")
    ax_max.set_title(f"{name}\n{proj_label}", fontsize=9)
    ax_max.axis("off")

    ax_roi.imshow(proj, cmap="gray")
    # Use a masked array so only ROI pixels get colored (zeros stay transparent)
    masked = np.ma.masked_where(binary_mask == 0, binary_mask)
    ax_roi.imshow(masked, cmap="hot", alpha=0.7, vmin=0, vmax=1)
    ax_roi.set_title(f"{name}\nROI Masks ({n_soma} soma)", fontsize=9)
    ax_roi.axis("off")

# Hide any unused axes (if n_planes is odd)
for i in range(n_planes, n_rows * N_COLS_PLANES):
    row      = i // N_COLS_PLANES
    col_base = (i %  N_COLS_PLANES) * 2
    axes[row, col_base].axis("off")
    axes[row, col_base + 1].axis("off")

fig.suptitle("All Imaging Planes — Max Projection and ROI Masks", fontsize=14)
plt.tight_layout()
plt.show()

# Choose which plane to inspect; change the index to select a different plane.
# The full list of available planes was printed by the Imaging Planes cell above.
PLANE = plane_names[0]

proc = nwb.processing[PLANE]
dff_series = proc["dff_timeseries"]["dff_timeseries"]
dff_data   = dff_series.data          # shape: (time, n_rois)
dff_ts     = np.array(dff_series.timestamps)

roi_table = proc["image_segmentation"]["roi_table"]
is_soma   = np.array(roi_table["is_soma"].data).astype(bool)
soma_indices = np.where(is_soma)[0]

print(f"Plane:                {PLANE}")
print(f"dF/F array shape:     {dff_data.shape}  (timepoints × ROIs)")
print(f"Session duration:     {dff_ts[-1] - dff_ts[0]:.1f} s")
print(f"Sampling rate:        {1 / np.median(np.diff(dff_ts)):.2f} Hz")
print(f"Soma ROIs available:  {len(soma_indices)}")

N_ROIS_PLOT = 5
SNIPPET_LEN = min(1000, dff_data.shape[0])

# NOTE: loads SNIPPET_LEN × n_soma_rois values from disk/network (soma columns only)
# Load a small initial snippet to rank ROIs by variance
snippet = np.array(dff_data[:SNIPPET_LEN, :][:, soma_indices])  # (SNIPPET_LEN, n_soma)
soma_variances = np.var(snippet, axis=0)  # variance already indexed to soma columns
top_soma_local  = np.argsort(soma_variances)[-N_ROIS_PLOT:][::-1]  # indices into soma_indices
top_roi_indices = soma_indices[top_soma_local]                       # indices into the full ROI table

print(f"Top {N_ROIS_PLOT} soma ROIs by variance (full ROI indices): {top_roi_indices}")

# Display 120 seconds of the session for these top ROIs
DISPLAY_WINDOW = (0, 120)  # seconds from session start

t0 = dff_ts[0]
time_mask  = (dff_ts - t0 >= DISPLAY_WINDOW[0]) & (dff_ts - t0 <= DISPLAY_WINDOW[1])
idx_start  = int(np.argmax(time_mask))
idx_end    = len(time_mask) - int(np.argmax(time_mask[::-1]))
t_plot     = dff_ts[idx_start:idx_end] - t0

fig, axes = plt.subplots(N_ROIS_PLOT, 1, figsize=(14, 2.5 * N_ROIS_PLOT), sharex=True)

for ax, roi_idx in zip(axes, top_roi_indices):
    trace = np.array(dff_data[idx_start:idx_end, int(roi_idx)])
    ax.plot(t_plot, trace, linewidth=0.6, color="steelblue")
    ax.set_ylabel(f"ROI {roi_idx}\n$\\Delta$F/F", fontsize=8)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

axes[-1].set_xlabel("Time from session start (s)")
fig.suptitle(f"{PLANE} — $\\Delta$F/F traces for top {N_ROIS_PLOT} soma ROIs (by variance)", y=1.01)
plt.tight_layout()
plt.show()

RASTER_WINDOW = (0, 600)   # seconds
RASTER_VMAX   = 0.5        # adjust to taste

mask = (dff_ts - dff_ts[0] >= RASTER_WINDOW[0]) & (dff_ts - dff_ts[0] <= RASTER_WINDOW[1])
i0, i1 = int(np.argmax(mask)), len(mask) - int(np.argmax(mask[::-1]))

# NOTE: loads ~RASTER_WINDOW duration × n_soma_rois values from disk/network
dff_soma = np.array(dff_data[i0:i1, soma_indices])  # (time, n_soma)
t_raster = dff_ts[i0:i1] - dff_ts[0]

fig, ax = plt.subplots(figsize=(14, 6))
im = ax.imshow(
    dff_soma.T,
    aspect="auto",
    extent=[t_raster[0], t_raster[-1], 0, dff_soma.shape[1]],
    cmap="inferno",
    vmin=0, vmax=RASTER_VMAX,
)
ax.set_xlabel("Time from session start (s)")
ax.set_ylabel("Soma ROI #")
ax.set_title(f"{PLANE} — $\\Delta$F/F raster ({len(soma_indices)} soma ROIs, first {RASTER_WINDOW[1]} s)")
plt.colorbar(im, ax=ax, label="$\\Delta$F/F")
plt.tight_layout()
plt.show()

# --- Running Speed ---
running_module = nwb.processing.get("running")
if running_module is not None:
    speed_key = "speed" if "speed" in running_module.keys() else "running_speed"
    running_ts_obj = running_module[speed_key]
    running_speed_data = np.array(running_ts_obj.data)
    running_speed_timestamps = np.array(running_ts_obj.timestamps)
    print(f"Running speed: {len(running_speed_data)} samples, "
          f"duration {running_speed_timestamps[-1] - running_speed_timestamps[0]:.1f} s")
else:
    print("Running speed not found in this NWB.")
    running_speed_data = np.array([])
    running_speed_timestamps = np.array([])

# --- Eye Tracking (pupil area) ---
pupil = nwb.processing['eye_tracking']['pupil']
eye_tracking_area       = np.array(pupil['area'])
eye_tracking_timestamps = np.array(pupil['timestamps'])
# Alternative tables in the same module (uncomment to use instead):
# corneal_reflection = nwb.processing['eye_tracking']['corneal_reflection']
# ellipse_fits       = nwb.processing['eye_tracking']['eye_tracked']

print(f"Eye tracking (pupil area): {len(eye_tracking_area)} samples, "
      f"duration {eye_tracking_timestamps[-1] - eye_tracking_timestamps[0]:.1f} s")

BEHAV_WINDOW = (0, 600)  # seconds to display

fig, axes = plt.subplots(3, 1, figsize=(14, 9), sharex=True)

# NOTE: each signal is zeroed to its *own* first timestamp, not a common session
# clock.  The three subplots therefore share the same x-axis range (0–600 s) but
# are on independent time axes and are not perfectly aligned with each other.
# See Section 11 for proper interpolation onto a single shared time axis.

# --- Running speed ---
rs_t    = running_speed_timestamps - running_speed_timestamps[0]
mask_rs = (rs_t >= BEHAV_WINDOW[0]) & (rs_t <= BEHAV_WINDOW[1])
axes[0].plot(rs_t[mask_rs], running_speed_data[mask_rs], lw=0.7, color="darkorange")
axes[0].set_ylabel("Running speed\n(cm/s)")
axes[0].set_title("Running Speed")
axes[0].spines["top"].set_visible(False)
axes[0].spines["right"].set_visible(False)

# --- Eye tracking (pupil area) ---
et_t    = eye_tracking_timestamps - eye_tracking_timestamps[0]
mask_et = (et_t >= BEHAV_WINDOW[0]) & (et_t <= BEHAV_WINDOW[1])
axes[1].plot(et_t[mask_et], eye_tracking_area[mask_et], lw=0.7, color="purple")
axes[1].set_ylabel("Pupil area\n(pixels²)")
axes[1].set_title("Eye Tracking (Pupil Area)")
axes[1].spines["top"].set_visible(False)
axes[1].spines["right"].set_visible(False)

# --- Mean ΔF/F across soma ROIs ---
t_dff_rel = dff_ts - dff_ts[0]
mask_dff  = (t_dff_rel >= BEHAV_WINDOW[0]) & (t_dff_rel <= BEHAV_WINDOW[1])
# NOTE: loads BEHAV_WINDOW duration × all ROIs from disk/network, then subsets to soma
mean_dff  = np.nanmean(np.array(dff_data[mask_dff, :])[:, soma_indices], axis=1)
axes[2].plot(t_dff_rel[mask_dff], mean_dff, lw=0.7, color="steelblue")
axes[2].set_ylabel("Mean $\\Delta$F/F\n(soma ROIs)")
axes[2].set_title(f"{PLANE} Mean $\\Delta$F/F")
axes[2].set_xlabel("Time from session start (s)")
axes[2].spines["top"].set_visible(False)
axes[2].spines["right"].set_visible(False)

fig.suptitle("Behavioral and Neural Overview (not time-aligned)", fontsize=13)
plt.tight_layout()
plt.show()

# List all available stimulus tables
stim_table_names = list(nwb.intervals.keys())
print("Stimulus tables in this NWB:")
for name in stim_table_names:
    tbl = nwb.intervals[name]
    print(f"  {name!r:55s}  ({len(tbl)} rows)")

# Inspect this session's oddball stimulus table
stim_name = 'Sensory-motor mismatch block_presentations'
stim_table = nwb.intervals[stim_name]

# Summarize stimulus types in this table
stim_df = stim_table.to_dataframe()

print(f"Table: {stim_name!r}")
print(f"Columns: {stim_table.colnames}")
print(f"TrialTypes: {set(stim_df['TrialType'])}")
stim_table[:10]

def extract_epochs(stim_name, stim_table, epochs):
    """Extract contiguous stimulus epochs from a presentation table.
    
    Scans the table row-by-row and groups consecutive rows with the same
    block / type value into a single epoch tuple:
        (stim_name, block_id, start_time, stop_time)
    """
    block_key = None
    for candidate in ["stim_block", "stimulus_block", "stimulus_type", "TrialType"]:
        if candidate in stim_table.colnames:
            block_key = candidate
            break

    if block_key is None:
        epoch_start = stim_table.start_time[0]
        epoch_stop  = stim_table.stop_time[-1]
        epochs.append((stim_name, stim_name, epoch_start, epoch_stop))
        return epochs

    epoch_start = stim_table.start_time[0]
    epoch_stop  = stim_table.stop_time[0]

    for i in range(len(stim_table)):
        this_block = stim_table[block_key][i]
        if i + 1 >= len(stim_table):
            epochs.append((stim_name, this_block, epoch_start, epoch_stop))
            break
        next_block = stim_table[block_key][i + 1]
        if next_block == this_block:
            epoch_stop = stim_table.stop_time[i + 1]
        else:
            epochs.append((stim_name, this_block, epoch_start, epoch_stop))
            epoch_start = stim_table.start_time[i + 1]
            epoch_stop  = stim_table.stop_time[i + 1]

    return epochs


epochs = []
for sn in stim_table_names:
    try:
        epochs = extract_epochs(sn, nwb.intervals[sn], epochs)
    except Exception:
        continue

epochs.sort(key=lambda x: x[2])
print(f"Total epochs found: {len(epochs)}")
for ep in epochs[:20]:
    print(f"  {ep[1]!s:<35}  {ep[2]:.1f} – {ep[3]:.1f} s  ({ep[3]-ep[2]:.1f} s)")

# Build color map: one color per stimulus table name
unique_stim_names = list({ep[0] for ep in epochs})

# Fixed color map — one visually distinct color per stimulus table.
# Mismatch/oddball tables use dark bold colors; paired controls use a lighter version.
STIM_COLORS = {
    # Oddball epochs (dark / bold)
    "Standard mismatch block_presentations":      "#4FC3CF",  # aquamarine
    "Sequence mismatch block_presentations":      "#B5541A",  # burnt orange
    "Duration mismatch block_presentations":      "#C8941A",  # goldenrod
    "Sensory-motor mismatch block_presentations": "#1B3A7A",  # navy blue
    # Paired control epochs (lighter version of the paired oddball)
    "Control block 1_presentations":              "#A8DDE6",  # light aquamarine (→ Standard)
    "Control block 2_presentations":              "#E0905E",  # light orange   (→ Sequence)
    "Control block 3_presentations":              "#EDD070",  # light goldenrod (→ Duration)
    "Control block 4_presentations":              "#7A9FC2",  # light navy     (→ Motor)
    # Other epochs
    "spontaneous_presentations":                  "#808080",  # grey
    "RF mapping_presentations":                   "#2E7D32",  # dark forest green
    "Zebra_presentations":                        "#8BC34A",  # lime / yellow-green
    "40 hz pulse train_presentations":            "#6A1B9A",  # deep indigo-purple
    "5 hz pulse train_presentations":             "#7B68EE",  # periwinkle/lavender
    "raised_cosine_presentations":                "#BF360C",  # deep brick red
}
# Fall back to neutral grey for any table not listed above.
color_map = {name: STIM_COLORS.get(name, "#999999") for name in unique_stim_names}

session_end = max(ep[3] for ep in epochs)

fig, ax = plt.subplots(figsize=(16, 2))

legend_handles = {}
for ep in epochs:
    stim_name_ep, block_id, t_start, t_stop = ep
    color = color_map[stim_name_ep]
    patch = ax.axvspan(t_start, t_stop, color=color, alpha=0.8)
    if stim_name_ep not in legend_handles:
        legend_handles[stim_name_ep] = patch

ax.set_xlim(0, session_end)
ax.set_ylim(0, 1)
ax.set_yticks([])
ax.set_xlabel("Time from session start (s)")
ax.set_title("Stimulus Epoch Structure")
ax.legend(
    legend_handles.values(), legend_handles.keys(),
    loc="upper right", bbox_to_anchor=(1, 2.5),
    fontsize=7, ncols=4, title="Stimulus table"
)
plt.tight_layout()
plt.show()

# Select the stimulus table — skip spontaneous activity tables
TYPE_COL = "TrialType"

non_spontaneous = [n for n in stim_table_names if "spontaneous" not in n.lower()]
STIM_TABLE_NAME = non_spontaneous[0]   # change if needed
stim_df = nwb.intervals[STIM_TABLE_NAME].to_dataframe()

print(f"Using stimulus table: {STIM_TABLE_NAME!r}")
print(f"Using type column:    {TYPE_COL!r}")
print(f"\nAvailable TrialTypes: {sorted(stim_df[TYPE_COL].unique().tolist())}")

# --- Choose which oddball type and standard type to align to ---
# Set these to match the TrialType values printed above.
ODDBALL_TYPE  = "halt"    # <-- EDIT THIS (e.g. 'halt', 'omission', 'orientation_45')
STANDARD_TYPE = "single"  # <-- EDIT THIS (the repeated baseline stimulus)

if ODDBALL_TYPE in stim_df[TYPE_COL].values:
    oddball_times = stim_df.loc[stim_df[TYPE_COL] == ODDBALL_TYPE, "start_time"].values
    print(f"Found {len(oddball_times)} presentations of '{ODDBALL_TYPE}'")
    print(f"First 10 onset times (s): {oddball_times[:10]}")
else:
    print(f"'{ODDBALL_TYPE}' not found. Available: {sorted(stim_df[TYPE_COL].unique().tolist())}")
    oddball_times = np.array([])

if STANDARD_TYPE in stim_df[TYPE_COL].values:
    standard_times = stim_df.loc[stim_df[TYPE_COL] == STANDARD_TYPE, "start_time"].values
    print(f"Found {len(standard_times)} '{STANDARD_TYPE}' presentations for comparison.")
else:
    print(f"'{STANDARD_TYPE}' not found. Available: {sorted(stim_df[TYPE_COL].unique().tolist())}")
    standard_times = np.array([])

INTERP_HZ = 100  # Hz for common time axis

# Build shared time axis spanning all signals
t_start_global = 0.0
t_end_global   = float(dff_ts[-1])
if len(running_speed_timestamps) > 0:
    t_end_global = max(t_end_global, float(running_speed_timestamps[-1]))
t_end_global = max(t_end_global, float(eye_tracking_timestamps[-1]))

time_axis = np.arange(t_start_global, t_end_global, step=1.0 / INTERP_HZ)

# Interpolate running speed
if len(running_speed_data) > 0:
    f_run = interpolate.interp1d(
        running_speed_timestamps, running_speed_data,
        kind="nearest", bounds_error=False, fill_value=0.0
    )
    interp_running = f_run(time_axis)
else:
    interp_running = np.zeros(len(time_axis))

# Interpolate pupil area
f_eye = interpolate.interp1d(
    eye_tracking_timestamps, eye_tracking_area,
    kind="nearest", bounds_error=False, fill_value=np.nan
)
interp_eye = f_eye(time_axis)

# NOTE: loads the full session × n_soma_rois into memory — this may take a moment
# for large files or slow network connections
f_dff = interpolate.interp1d(
    dff_ts,
    np.array(dff_data[:, soma_indices]),  # shape (time, n_soma)
    axis=0, kind="nearest", bounds_error=False, fill_value=0.0
)
interp_dff = f_dff(time_axis).T  # shape (n_soma, time)

print(f"Common time axis: {len(time_axis)} samples at {INTERP_HZ} Hz ({time_axis[-1]:.1f} s)")
print(f"interp_running shape: {interp_running.shape}")
print(f"interp_eye shape:     {interp_eye.shape}")
print(f"interp_dff shape:     {interp_dff.shape}  (soma ROIs × time)")

VIEW_START = 0                    # seconds
VIEW_END   = int(time_axis[-1])  # full session; set to e.g. 600 to zoom in on a region

s0 = int(VIEW_START * INTERP_HZ)
s1 = int(VIEW_END   * INTERP_HZ)

t_view       = time_axis[s0:s1]
running_view = interp_running[s0:s1]
eye_view     = interp_eye[s0:s1]
dff_view     = interp_dff[:, s0:s1]  # (n_soma, time)


def draw_epoch_bands(ax, epochs, view_start, view_end, cmap=None, alpha=0.30):
    """Shade each epoch as a semi-transparent colored band on the given axis.

    Uses ``ax.axvspan`` so the bands span the full y-axis without capturing
    the current ylim and without interfering with autoscaling.

    Parameters
    ----------
    cmap : dict, optional
        A {label: color} mapping.  When provided (e.g. the ``color_map`` dict
        produced by the Epoch Structure plot) the same colors are reused so
        both figures are consistent.  When omitted a local palette is built.
    alpha : float
        Opacity of the shaded bands (0 = invisible, 1 = solid).
    """
    if cmap is None:
        unique_labels = list({ep[0] for ep in epochs})
        tab_colors    = plt.cm.tab20(np.linspace(0, 1, len(unique_labels)))
        cmap          = dict(zip(unique_labels, tab_colors))

    legend_handles = {}
    for ep in epochs:
        stim_name_ep, _, t_s, t_e = ep
        if t_e < view_start or t_s > view_end:
            continue
        color = cmap.get(stim_name_ep, "gray")
        patch = ax.axvspan(
            max(t_s, view_start), min(t_e, view_end),
            color=color, alpha=alpha, zorder=0
        )
        legend_handles[stim_name_ep] = patch
    return legend_handles


# GridSpec: 3 rows × 2 cols. Column 0 is the plots; column 1 is a narrow colorbar
# column only used by the dF/F row — keeps all plot panels the same width.
height_ratios = [1.2, 1.2, 3.5]
fig = plt.figure(figsize=(15, 10))
gs  = gridspec.GridSpec(3, 2, figure=fig,
                        height_ratios=height_ratios,
                        width_ratios=[40, 1])

ax0 = fig.add_subplot(gs[0, 0])
ax1 = fig.add_subplot(gs[1, 0], sharex=ax0)
ax2 = fig.add_subplot(gs[2, 0], sharex=ax0)
cax = fig.add_subplot(gs[2, 1])   # colorbar column, dF/F row only
axes = np.array([ax0, ax1, ax2])

# --- Running speed ---
ax0.plot(t_view, running_view, lw=0.7, color="darkorange")
ax0.set_ylabel("Running\nspeed (cm/s)")
ax0.set_title("Running Speed")
ax0.spines["top"].set_visible(False); ax0.spines["right"].set_visible(False)
lh = draw_epoch_bands(ax0, epochs, VIEW_START, VIEW_END, cmap=color_map)

# --- Eye area ---
ax1.plot(t_view, eye_view, lw=0.7, color="purple")
ax1.set_ylabel("Pupil area\n(pixels²)")
ax1.set_title("Eye Tracking (Pupil Area)")
ax1.spines["top"].set_visible(False); ax1.spines["right"].set_visible(False)
draw_epoch_bands(ax1, epochs, VIEW_START, VIEW_END, cmap=color_map)

# --- ΔF/F raster ---
im = ax2.imshow(
    dff_view,
    aspect="auto",
    extent=[t_view[0], t_view[-1], 0, dff_view.shape[0]],
    cmap="inferno", vmin=0, vmax=0.5, origin="lower"
)
ax2.set_ylabel("Soma ROI #")
ax2.set_xlabel("Time (s)")
ax2.set_title(f"{PLANE} — $\\Delta$F/F (all soma ROIs)")
fig.colorbar(im, cax=cax, label="$\\Delta$F/F")

if lh:
    ax0.legend(
        lh.values(), lh.keys(),
        ncols=4, fontsize=7,
        loc="upper left", bbox_to_anchor=(0, 2.7),
        title="Stimulus table"
    )

fig.suptitle(
    f"Time-Aligned Modalities — Predictive Processing Session\n"
    f"({VIEW_START}–{VIEW_END} s, plane {PLANE})",
    fontsize=13
)
plt.tight_layout()
plt.show()

WINDOW_PRE  = -1.0  # seconds before stimulus onset
WINDOW_POST =  2.0  # seconds after stimulus onset

window_len = int((WINDOW_POST - WINDOW_PRE) * INTERP_HZ)
window_t   = np.linspace(WINDOW_PRE, WINDOW_POST, window_len)


def get_stim_windows(signal, stim_times, interp_hz, window_pre, window_post):
    """Extract fixed-length windows around each stimulus onset.

    Parameters
    ----------
    signal : ndarray, shape (n_timepoints,) or (n_features, n_timepoints)
        Data on the global uniform time axis.
    stim_times : array-like
        Stimulus onset times in seconds.
    interp_hz : float
        Sampling rate of the uniform time axis that ``signal`` has already been
        interpolated onto (e.g. ``INTERP_HZ``).  The function uses this rate to
        convert onset times to sample indices, so ``signal`` must be on a
        regularly-sampled grid at exactly this rate.
    window_pre, window_post : float
        Window boundaries relative to onset (window_pre < 0 < window_post).

    Returns
    -------
    ndarray, shape (n_trials, window_len) or (n_features, n_trials, window_len)
    """
    wlen = int((window_post - window_pre) * interp_hz)
    n_t  = signal.shape[-1]
    slices = []
    for t in stim_times:
        i0 = int((t + window_pre) * interp_hz)
        i1 = i0 + wlen
        if 0 <= i0 and i1 <= n_t:
            slices.append(signal[..., i0:i1])
    if not slices:
        raise ValueError("No valid windows — check stim_times and window bounds.")
    if signal.ndim == 1:
        return np.array(slices)          # (n_trials, window_len)
    return np.stack(slices, axis=1)      # (n_features, n_trials, window_len)


# Extract windows for oddball stimuli
run_win_odd = get_stim_windows(interp_running, oddball_times,  INTERP_HZ, WINDOW_PRE, WINDOW_POST)
run_win_std = get_stim_windows(interp_running, standard_times, INTERP_HZ, WINDOW_PRE, WINDOW_POST)

eye_win_odd = get_stim_windows(interp_eye, oddball_times,  INTERP_HZ, WINDOW_PRE, WINDOW_POST)
eye_win_std = get_stim_windows(interp_eye, standard_times, INTERP_HZ, WINDOW_PRE, WINDOW_POST)

dff_win_odd = get_stim_windows(interp_dff, oddball_times,  INTERP_HZ, WINDOW_PRE, WINDOW_POST)
dff_win_std = get_stim_windows(interp_dff, standard_times, INTERP_HZ, WINDOW_PRE, WINDOW_POST)
# dff_win_odd shape: (n_soma, n_trials, window_len)

print(f"Oddball windows:  {run_win_odd.shape[0]} trials")
print(f"Standard windows: {run_win_std.shape[0]} trials")
print(f"dF/F window shape (soma × trials × time): {dff_win_odd.shape}")

N_PLOT_CELLS = 20  # number of soma ROIs to show (capped at available)

n_show = min(N_PLOT_CELLS, dff_win_odd.shape[0])

# Shared color scale: 99th-percentile ceiling across both conditions so the
# two columns are directly comparable.
vmax = float(np.nanpercentile(
    np.concatenate([dff_win_odd[:n_show].ravel(), dff_win_std[:n_show].ravel()]), 99
))

fig, axes = plt.subplots(
    n_show, 2,
    figsize=(11, n_show * 1.5),
    sharex=True,
)

for i in range(n_show):
    for col, (wins, label, title_color) in enumerate([
        (dff_win_std[i], "Standard",                   "steelblue"),
        (dff_win_odd[i], f"Oddball ({ODDBALL_TYPE})",  "tomato"),
    ]):
        ax = axes[i, col]
        im = ax.imshow(
            wins,
            aspect="auto",
            extent=[window_t[0], window_t[-1], 0, wins.shape[0]],
            cmap="inferno", vmin=0.0, vmax=vmax,
            origin="lower", interpolation="nearest",
        )
        ax.axvline(0, color="white", lw=0.8, ls="--", alpha=0.7)
        ax.set_yticks([])
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        if i == 0:
            ax.set_title(f"{label}  (n={wins.shape[0]})", fontsize=10,
                         color=title_color, fontweight="bold", pad=4)
        if i < n_show - 1:
            ax.tick_params(labelbottom=False)

    axes[i, 0].set_ylabel(f"ROI {soma_indices[i]}", fontsize=7,
                          rotation=0, labelpad=36, va="center")

axes[-1, 0].set_xlabel("Time rel. to onset (s)")
axes[-1, 1].set_xlabel("Time rel. to onset (s)")

# Single shared colorbar pinned to the right of the figure
fig.subplots_adjust(right=0.87, hspace=0.05, wspace=0.04)
cbar_ax = fig.add_axes([0.90, 0.15, 0.018, 0.70])
fig.colorbar(im, cax=cbar_ax, label="ΔF/F")

fig.suptitle(
    f"Trial × Time Windows — First {n_show} Soma ROIs\n"
    f"(Oddball: {ODDBALL_TYPE!r},  plane: {PLANE})",
    fontsize=12,
)
plt.show()

def mean_sem(windows):
    """Compute mean and SEM across trials. ``windows`` must have shape (n_trials, time)."""
    n = windows.shape[0]
    m = np.nanmean(windows, axis=0)
    s = np.nanstd(windows, axis=0) / np.sqrt(n)
    return m, s


run_mean_odd, run_sem_odd = mean_sem(run_win_odd)
run_mean_std, run_sem_std = mean_sem(run_win_std)

eye_mean_odd, eye_sem_odd = mean_sem(eye_win_odd)
eye_mean_std, eye_sem_std = mean_sem(eye_win_std)

# For dF/F: first average across ROIs within each trial, then average trials
# dff_win_odd: (n_soma, n_trials, time) → average over soma → (n_trials, time)
dff_mean_per_trial_odd = np.nanmean(dff_win_odd, axis=0)  # (n_trials, time)
dff_mean_per_trial_std = np.nanmean(dff_win_std, axis=0)
dff_mean_odd, dff_sem_odd = mean_sem(dff_mean_per_trial_odd)
dff_mean_std, dff_sem_std = mean_sem(dff_mean_per_trial_std)

print("Average response shapes computed.")

fig, axes = plt.subplots(3, 1, figsize=(10, 11), sharex=True,
                         gridspec_kw={"height_ratios": [1, 1, 1.6]})

# --- Panel 1: Running speed ---
axes[0].axvline(0, color="gray", lw=0.8, ls="--", alpha=0.7)
axes[0].fill_between(window_t, run_mean_odd - run_sem_odd,
                               run_mean_odd + run_sem_odd, alpha=0.25, color="tomato")
axes[0].fill_between(window_t, run_mean_std - run_sem_std,
                               run_mean_std + run_sem_std, alpha=0.25, color="steelblue")
axes[0].plot(window_t, run_mean_odd, color="tomato",   lw=1.5, label=f"Oddball ({run_win_odd.shape[0]} trials)")
axes[0].plot(window_t, run_mean_std, color="steelblue", lw=1.5, label=f"Standard ({run_win_std.shape[0]} trials)")
axes[0].set_ylabel("Running speed\n(cm/s)")
axes[0].set_title("Stimulus-Aligned Running Speed")
axes[0].legend(loc="upper right", fontsize=9)
axes[0].spines["top"].set_visible(False); axes[0].spines["right"].set_visible(False)

# --- Panel 2: Pupil area ---
axes[1].axvline(0, color="gray", lw=0.8, ls="--", alpha=0.7)
axes[1].fill_between(window_t, eye_mean_odd - eye_sem_odd,
                               eye_mean_odd + eye_sem_odd, alpha=0.25, color="tomato")
axes[1].fill_between(window_t, eye_mean_std - eye_sem_std,
                               eye_mean_std + eye_sem_std, alpha=0.25, color="steelblue")
axes[1].plot(window_t, eye_mean_odd, color="tomato",   lw=1.5, label="Oddball")
axes[1].plot(window_t, eye_mean_std, color="steelblue", lw=1.5, label="Standard")
axes[1].set_ylabel("Pupil area\n(pixels²)")
axes[1].set_title("Stimulus-Aligned Pupil Area")
axes[1].legend(loc="upper right", fontsize=9)
axes[1].spines["top"].set_visible(False); axes[1].spines["right"].set_visible(False)

# --- Panel 3: Mean ΔF/F ---
axes[2].axvline(0, color="gray", lw=0.8, ls="--", alpha=0.7)
axes[2].fill_between(window_t, dff_mean_odd - dff_sem_odd,
                               dff_mean_odd + dff_sem_odd, alpha=0.25, color="tomato")
axes[2].fill_between(window_t, dff_mean_std - dff_sem_std,
                               dff_mean_std + dff_sem_std, alpha=0.25, color="steelblue")
axes[2].plot(window_t, dff_mean_odd, color="tomato",   lw=2.0, label=f"Oddball")
axes[2].plot(window_t, dff_mean_std, color="steelblue", lw=2.0, label="Standard")
axes[2].axhline(0, color="black", lw=0.5, ls=":")
axes[2].set_ylabel("Mean $\\Delta$F/F\n(soma ROIs)")
axes[2].set_xlabel("Time relative to stimulus onset (s)")
axes[2].set_title(f"{PLANE} — Population-Averaged Neural Mismatch Response")
axes[2].legend(loc="upper right", fontsize=9)
axes[2].spines["top"].set_visible(False); axes[2].spines["right"].set_visible(False)

fig.suptitle(
    f"Stimulus-Aligned Average Response: Oddball vs. Standard\n"
    f"(Oddball type: {ODDBALL_TYPE!r}, plane: {PLANE})",
    fontsize=13
)
plt.tight_layout()
plt.show()

# Per-ROI average response: (n_soma, time)
dff_roi_avg_odd = np.nanmean(dff_win_odd, axis=1)  # average over trials
dff_roi_avg_std = np.nanmean(dff_win_std, axis=1)

# Difference (mismatch response)
mismatch_response = dff_roi_avg_odd - dff_roi_avg_std

# Sort by time of peak absolute response after onset
onset_idx  = int(-WINDOW_PRE * INTERP_HZ)
peak_times = np.argmax(np.abs(mismatch_response[:, onset_idx:]), axis=1)
sort_order = np.argsort(peak_times)

fig, axes = plt.subplots(1, 3, figsize=(18, 7))

for ax, data, title, cmap, vlim, cbar_label in [
    (axes[0], dff_roi_avg_std[sort_order], "Standard",
     "inferno", (0, 0.3),    "Mean $\\Delta$F/F (standard)"),
    (axes[1], dff_roi_avg_odd[sort_order], f"Oddball: {ODDBALL_TYPE}",
     "inferno", (0, 0.3),    "Mean $\\Delta$F/F (oddball)"),
    (axes[2], mismatch_response[sort_order], "Response Difference (Oddball $-$ Standard)",
     "RdBu_r",  (-0.15, 0.15), "$\\Delta\\Delta$F/F (oddball $-$ standard)"),
]:
    im = ax.imshow(
        data,
        aspect="auto",
        extent=[window_t[0], window_t[-1], 0, data.shape[0]],
        cmap=cmap, vmin=vlim[0], vmax=vlim[1], origin="lower"
    )
    ax.axvline(0, color="white", lw=1, ls="--", alpha=0.8)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Soma ROI # (sorted by peak time)")
    ax.set_title(title)
    plt.colorbar(im, ax=ax, label=cbar_label, shrink=0.6)

fig.suptitle(
    f"Per-ROI Stimulus-Aligned Responses — {PLANE}\n"
    f"Sorted by peak response time after stimulus onset",
    fontsize=13
)
plt.tight_layout()
plt.show()