OpenScope’s Sequence Learning Dataset

import warnings
warnings.filterwarnings('ignore')

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

import os

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from mpl_interactions import hyperslicer
from scipy import interpolate
from scipy.stats import ttest_ind

%matplotlib inline

session_files = pd.read_csv("../../data/seqlearn_sessions.csv")
session_files

# Function to convert string representation of a set to an actual set
def parse_set_string(set_string):
    # Remove curly braces, split by commas, strip spaces and single quotes
    return set(item.strip("'") for item in set_string.strip("{}").split(", "))

subjects = session_files.groupby('sub_name').agg(
    n_sessions=('session_id', 'nunique'),  # Count unique session IDs
    stim_types=('stim_types', lambda x: set().union(*x.apply(parse_set_string))),  # Union of all sets of stim_types
    genotype=('sub_genotype', 'unique')
).reset_index()

subjects

n_sessions = len(session_files["session_id"].value_counts())
subjects_info = session_files.groupby(["specimen_name", "sub_sex"]).size().reset_index().to_dict()
m_count = len([sex for sex in subjects_info["sub_sex"].values() if sex == "M"])
f_count = len([sex for sex in subjects_info["sub_sex"].values() if sex == "F"])

print("Dandiset Overview:")
print(len(session_files), "files")
print(len(subjects_info["specimen_name"]), "subjects", m_count, "males", f_count,"females")

dandiset_id = "000617"
sequence_dandi_filepath = "sub-688425/sub-688425_ses-1306855381-acq-1307046775_ophys.nwb"
random_dandi_filepath = "sub-688425/sub-688425_ses-1307534620-acq-1307650789_ophys.nwb"
download_loc = "."

# This can sometimes take a while depending on the size of the file
seq_io = dandi_download_open(dandiset_id, sequence_dandi_filepath, download_loc)
seq_nwb = seq_io.read()

# This can sometimes take a while depending on the size of the file
rand_io = dandi_download_open(dandiset_id, random_dandi_filepath, download_loc)
rand_nwb = rand_io.read()

# view the contents of the nwb interactively
seq_nwb

print("Subject ID",seq_nwb.subject.subject_id)
print('Imaging Depth', seq_nwb.lab_meta_data['metadata'].imaging_depth)
print('Location', seq_nwb.imaging_planes['imaging_plane_1'].location)

print("Subject ID",rand_nwb.subject.subject_id)
print('Imaging Depth', rand_nwb.lab_meta_data['metadata'].imaging_depth)
print('Location', rand_nwb.imaging_planes['imaging_plane_1'].location)

seq_nwb.stimulus_template.keys()

n_cols = 2
n_rows = 2
fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols*2, n_rows*2))

if len(axes.shape) == 1:
    axes = axes.reshape((1, axes.shape[0]))

for i, template_name in enumerate(seq_nwb.stimulus_template.keys()):
    template_img = seq_nwb.stimulus_template[template_name].data[:,:,0]
    template_img = np.rot90(template_img, k=3)

    ax_row = int(i / n_cols)
    ax_col = i % n_cols
    axes[ax_row][ax_col].imshow(template_img, cmap="gray")
    axes[ax_row][ax_col].set_title(template_name, fontsize=8)

for ax in axes.flat:
    ax.xaxis.set_ticks([])
    ax.yaxis.set_ticks([])

fig.tight_layout()

%matplotlib ipympl
plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)

# change this key to view other stim template movies
template_key = "movie_clip_B"
template = seq_nwb.stimulus_template[template_key].data
template = np.transpose(template)
view = hyperslicer(template, play_buttons=True, cmap="gray")

print(seq_nwb.intervals.keys())
print(rand_nwb.intervals.keys())

seq_stim_table = seq_nwb.intervals["movie_clip_B_presentations"]
seq_stim_table[:10]

print("Frame nums:", set(seq_stim_table.frame))

seq_stim_select = lambda row: row.frame.item() == 0.0
seq_stim_times = [float(seq_stim_table[i].start_time) for i in range(len(seq_stim_table)) if seq_stim_select(seq_stim_table[i])]
print("Selected stim times:", len(seq_stim_times))

rand_stim_table = rand_nwb.intervals["movie_clip_B_presentations"]
rand_stim_table[:10]

print("Frame nums:", set(rand_stim_table.frame))

rand_stim_select = lambda row: row.frame.item() == 0.0
rand_stim_times = [float(rand_stim_table[i].start_time) for i in range(len(rand_stim_table)) if rand_stim_select(rand_stim_table[i])]
print("Selected stim times:", len(rand_stim_times))

%matplotlib inline
fig, axes = plt.subplots(1,2, figsize=(10,30))
axes[0].imshow(rand_nwb.processing['ophys']['images']['max_projection'])
axes[1].imshow(rand_nwb.processing['ophys']['images']['segmentation_mask_image'])

seq_dff = seq_nwb.processing["ophys"]["dff"]
seq_dff_trace = np.array(seq_dff.roi_response_series["traces"].data)
seq_dff_timestamps = np.array(seq_dff.roi_response_series["traces"].timestamps)

print(seq_dff_trace.shape)
print(seq_dff_timestamps.shape)

seq_avg_dff_trace = np.average(seq_dff_trace, axis=1)

rand_dff = rand_nwb.processing["ophys"]["dff"]
rand_dff_trace = np.array(rand_dff.roi_response_series["traces"].data)
rand_dff_timestamps = np.array(rand_dff.roi_response_series["traces"].timestamps)

print(rand_dff_trace.shape)
print(rand_dff_timestamps.shape)

rand_avg_dff_trace = np.average(rand_dff_trace, axis=1)

%matplotlib inline
n_rois = seq_dff_trace.shape[1]
plt.imshow(seq_dff_trace.transpose(), extent=[seq_dff_timestamps[0], seq_dff_timestamps[-1], 0, n_rois], aspect='auto', vmin=0, vmax=1, interpolation='None')
# plt.yticks(np.arange(n_rois)+0.5, np.arange(n_rois))
plt.ylabel("ROI #")
plt.xlabel("time (s)")
plt.title("DF/F for all ROIs throughout entire session")
cbar = plt.colorbar()
cbar.set_label('DF/F (%)')

seq_running = seq_nwb.processing["running"]["speed"]
seq_running_trace = np.array(seq_running.data)
seq_running_timestamps = np.array(seq_running.timestamps)

print(seq_running_trace.shape)
print(seq_running_timestamps.shape)

rand_running = rand_nwb.processing["running"]["speed"]
rand_running_trace = np.array(rand_running.data)
rand_running_timestamps = np.array(rand_running.timestamps)

print(rand_running_trace.shape)
print(rand_running_timestamps.shape)

plt.plot(seq_running_timestamps, seq_running_trace)
plt.xlabel("Time (s)")
plt.ylabel("Speed (cm/s)")
plt.title("Mouse running throughout session")

# extract epoch times from stim table where stimulus rows have a different 'block' than following row
# returns list of epochs, where an epoch is of the form (stimulus name, stimulus block, start time, stop time)
def extract_epochs_from_table(stim_name, stim_table, epochs):
    
    # specify a current epoch stop and start time
    epoch_start = stim_table.start_time[0]
    epoch_stop = stim_table.stop_time[0]

    # for each row, try to extend current epoch stop_time
    for i in range(len(stim_table)):
        this_block = stim_table.stimulus_block[i]
        # if end of table, end the current epoch
        if i+1 >= len(stim_table):
            epochs.append((stim_name, this_block, epoch_start, epoch_stop))
            break
            
        next_block = stim_table.stimulus_block[i+1]
        # if next row is the same stim block, push back epoch_stop time
        if next_block == this_block:
            epoch_stop = stim_table.stop_time[i+1]
        # otherwise, end the current epoch, start new epoch
        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

def extract_all_epochs(nwb):
    # extract epochs from all valid stimulus tables
    epochs = []
    for stim_name in nwb.intervals.keys():
        stim_table = nwb.intervals[stim_name]
        try:
            epochs = extract_epochs_from_table(stim_name, stim_table, epochs)
        except:
            continue
    return epochs

%matplotlib inline
### make plot of chosen trace over time with colored epoch sections

def plot_trace_over_epochs(trace_arrs, timestamps_arrs, epochs, time_start=None, time_end=None, title=None, trace_labels=None, yaxlabels=None, xlabel=None):
    assert len(trace_arrs) == len(timestamps_arrs), "there must be an equal number of traces and timestamps arrays"
    if trace_labels is not None:
        assert len(trace_arrs) == len(trace_labels), "there must be an equal number of traces and trace labels arrays"
    if yaxlabels is not None:
        assert len(trace_arrs) == len(yaxlabels),  "there must be an equal number of traces and y-axis labels arrays"

    fig, ax = plt.subplots(figsize=(15,5), dpi=300)
    if time_start is None:
        time_start = np.min(np.concatenate(timestamps_arrs))
    if time_end is None:
        time_end = np.max(np.concatenate(timestamps_arrs))

    # filter epochs which aren't at least partially in the time window
    bounded_epochs = {epoch for epoch in epochs if epoch[2] < time_end and epoch[3] > time_start}

    # assign unique color to each stimulus name
    stim_names = list({epoch[0] for epoch in bounded_epochs})
    colors = plt.cm.rainbow(np.linspace(0,1,len(stim_names)))
    stim_color_map = {stim_names[i]:colors[i] for i in range(len(stim_names))}

    key = {}
    
    y_hi = np.max(np.concatenate(trace_arrs)) # change these to manually set height of the plot
    y_lo = np.min(np.concatenate(trace_arrs))
    # draw colored rectangles for each epoch
    for epoch in bounded_epochs:
        stim_name, stim_block, epoch_start, epoch_end = epoch
        color = stim_color_map[stim_name]
        rec = ax.add_patch(mpl.patches.Rectangle((epoch_start, y_lo), epoch_end-epoch_start, 100, alpha=0.3, facecolor=color, antialiased=False))
        key[(stim_name)] = rec

    ax.set_xlim(time_start, time_end)
    ax.set_ylim(y_lo, y_hi)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if title is not None:
        ax.set_title(title)
    if yaxlabels is not None:
        ax.set_ylabel("\n".join(yaxlabels))

    for i in range(len(trace_arrs)):
        # next_color = plt.rcParams["axes.prop_cycle"].by_key()["color"][i]
        # this_ax = ax.twinx()
        line = ax.plot(timestamps_arrs[i], trace_arrs[i])[0]
        if trace_labels is not None:
            key[(trace_labels[i])] = line

    print(key)
    fig.legend(key.values(), key.keys(), loc="lower right", bbox_to_anchor=(1.25, 0.25))
    
    return ax

seq_epochs = extract_all_epochs(seq_nwb)

# epochs take the form (stimulus name, stimulus block, start time, stop time)
print("Num epochs:",len(seq_epochs))
seq_epochs.sort(key=lambda x: x[2])
for epoch in seq_epochs:
    print(epoch)

# can set these manually to get a closer look at the timeline
time_start = min(seq_epochs, key=lambda epoch: epoch[1])[1]
time_end = max(seq_epochs, key=lambda epoch:epoch[2])[2]
# time_start = 3000
# time_end = 3100

# can set this to change what trace is displayed alongside epochs
display_trace = seq_avg_dff_trace * 100 # to yield percentage
# unit_idx = 30
# display_trace = dff_trace[:,unit_idx] * 100

# adjust these to zoom in to a narrow slice of time
start, stop = None, None
ax = plot_trace_over_epochs([display_trace, seq_running_trace], [seq_dff_timestamps, seq_running_timestamps], seq_epochs, start, stop, "Average DFF and running during a period of sequential session", ['avg dff', 'running'], ["Average DF/F (%)","Running Speed (cm/s)"], "Time (s)")
plt.tight_layout()
plt.show()

# adjust these to zoom in to a narrow slice of time
start, stop = 1000, 1200
ax = plot_trace_over_epochs([display_trace, seq_running_trace], [seq_dff_timestamps, seq_running_timestamps], seq_epochs, start, stop, "Average DFF and running throughout sequential session", ['avg dff', 'running'], ["Average DF/F (%)","Running Speed (cm/s)"], "Time (s)")
plt.tight_layout()
plt.show()

rand_epochs = extract_all_epochs(rand_nwb)

# epochs take the form (stimulus name, stimulus block, start time, stop time)
print("Num epochs:",len(rand_epochs))
rand_epochs.sort(key=lambda x: x[2])
# for epoch in rand_epochs:
#     print(epoch)

# can set these manually to get a closer look at the timeline
time_start = min(rand_epochs, key=lambda epoch: epoch[1])[1]
time_end = max(rand_epochs, key=lambda epoch:epoch[2])[2]
# time_start = 3000
# time_end = 3100

# can set this to change what trace is displayed alongside epochs
display_trace = rand_avg_dff_trace * 100 # to yield percentage
# unit_idx = 30
# display_trace = dff_trace[:,unit_idx] * 100

# adjust these to zoom in to a narrow slice of time
start, stop = None, None
ax = plot_trace_over_epochs([display_trace, rand_running_trace], [rand_dff_timestamps, rand_running_timestamps], rand_epochs, start, stop, "Average DFF and running throughout random session", ['avg dff', 'running'], ["Average DF/F (%)","Running Speed (cm/s)"], "Time (s)")
plt.tight_layout()
plt.show()

window_start_time = -2
window_end_time = 3
interp_hz = 10

# generate regularly-space x values and interpolate along it
time_axis = np.arange(seq_dff_timestamps[0], seq_dff_timestamps[-1], step=(1/interp_hz))
interp_dff = []

# interpolate channel by channel to save RAM
for channel in range(seq_dff_trace.shape[1]):
    f = interpolate.interp1d(seq_dff_timestamps, seq_dff_trace[:,channel], axis=0, kind="nearest", fill_value="extrapolate")
    interp_dff.append(f(time_axis))

interp_dff = np.array(interp_dff)

print(interp_dff.shape)

# validate window bounds
if window_start_time > 0:
    raise ValueError("start time must be non-positive number")
if window_end_time <= 0:
    raise ValueError("end time must be positive number")
    
# get event windows
windows = []
window_length = int((window_end_time-window_start_time) * interp_hz)

for stim_ts in seq_stim_times:
    # convert time to index
    start_idx = int( (stim_ts + window_start_time - seq_dff_timestamps[0]) * interp_hz )
    end_idx = start_idx + window_length
 
    # bounds checking
    if start_idx < 0 or end_idx > interp_dff.shape[1]:
        continue
        
    windows.append(interp_dff[:,start_idx:end_idx])
    
if len(windows) == 0:
    raise ValueError("There are no windows for these timestamps")

windows = np.array(windows) * 100 # x100 to convert values to dF/F percentage
neuronwise_windows = np.swapaxes(windows,0,1)

print(neuronwise_windows.shape)

%matplotlib inline
def show_dff_response(ax, dff, window_start_time, window_end_time, aspect="auto", vmin=None, vmax=None, yticklabels=[], skipticks=1, xlabel="Time (s)", ylabel="ROI", cbar=True, cbar_label=None):
    if len(dff) == 0:
        print("Input data has length 0; Nothing to display")
        return

    img = ax.imshow(dff, aspect=aspect, extent=[window_start_time, window_end_time, 0, len(dff)], vmin=vmin, vmax=vmax)
    if cbar:
        ax.colorbar(img, shrink=0.5, label=cbar_label)

    ax.plot([0,0],[0, len(dff)], ":", color="white", linewidth=1.0)

    if len(yticklabels) != 0:
        ax.set_yticks(range(len(yticklabels)))
        ax.set_yticklabels(yticklabels, fontsize=8)

        n_ticks = len(yticklabels[::skipticks])
        ax.yaxis.set_major_locator(plt.MaxNLocator(n_ticks))

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

def show_many_responses(windows, rows, cols, window_idxs=None, title=None, subplot_title="", xlabel=None, ylabel=None, cbar_label=None, vmin=0, vmax=100):
    if window_idxs is None:
        window_idxs = range(len(windows))
    windows = windows[window_idxs]
    
    # handle case with no input data
    if len(windows) == 0:
        print("Input data has length 0; Nothing to display")
        return
    # handle cases when there aren't enough windows for number of rows
    if len(windows) < rows*cols:
        rows = (len(windows) // cols) + 1

    fig, axes = plt.subplots(rows, cols, figsize=(2*cols+2, 2*rows+2), layout="constrained")
    # handle case when there's only one row
    if len(axes.shape) == 1:
        axes = axes.reshape((1, axes.shape[0]))
    
    for i in range(rows*cols):
        ax_row = int(i // cols)
        ax_col = i % cols
        ax = axes[ax_row][ax_col]
        
        if i > len(windows)-1:
            ax.set_visible(False)
            continue

        window = windows[i]
        show_dff_response(ax, window, window_start_time, window_end_time, xlabel=xlabel, ylabel=ylabel, cbar=False, vmin=vmin, vmax=vmax)
        ax.set_title(f"{subplot_title} {window_idxs[i]}")
        if ax_row != rows-1:
            ax.get_xaxis().set_visible(False)
        if ax_col != 0:
            ax.get_yaxis().set_visible(False)

    fig.suptitle(title)
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    colorbar = fig.colorbar(mpl.cm.ScalarMappable(norm=norm), ax=axes, shrink=1.5/rows, label=cbar_label)

show_many_responses(neuronwise_windows,
                    6,
                    15,
                    title="Response of Some ROIs across trials",
                    subplot_title="ROI",
                    xlabel="time (s)",
                    ylabel="trial",
                    cbar_label="$\Delta$F/F (%)")

# get the index within the window that stimulus occurs (time 0)
stimulus_onset_idx = int(-window_start_time * interp_hz)
baseline = windows[:,:,0:stimulus_onset_idx]
evoked_responses = windows[:,:,stimulus_onset_idx:]

print(stimulus_onset_idx)
print(baseline.shape)
print(evoked_responses.shape)

mean_trial_responses = np.mean(evoked_responses, axis=2)
mean_trial_baselines = np.mean(baseline, axis=2)

n = mean_trial_responses.shape[0]
t,p = ttest_ind(mean_trial_responses, mean_trial_baselines)
IC3_selected_rois = np.where(p < 0.05 / n)[0]
print(f"Selected ROIs {IC3_selected_rois}")

show_many_responses(neuronwise_windows[IC3_selected_rois],
                    6,
                    15,
                    title="Response of Some ROIs across trials",
                    subplot_title="ROI",
                    xlabel="time (s)",
                    ylabel="trial",
                    cbar_label="$\Delta$F/F (%)")

dandi_movie_filepath = "sub-688425/sub-688425_ses-1306855381-acq-1307046775-raw-movies_ophys.nwb"

# This can sometimes take a while depending on the size of the file
io = dandi_download_open(dandiset_id, dandi_movie_filepath, download_loc)
nwb = io.read()

# start_time = flr_timestamps[0]
start_time = 540
# end_time = flr_timestamps[-1]
end_time = 600

start_idx, end_idx = np.searchsorted(seq_dff_timestamps, [start_time, end_time])
print(start_idx)
print(end_idx)

raw_movie = nwb.acquisition["raw_suite2p_motion_corrected"].data
flr_timestamps = np.array(seq_dff.roi_response_series["traces"].timestamps)

print(raw_movie.shape)
print(flr_timestamps.shape)

%matplotlib ipympl
plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
view = hyperslicer(raw_movie[start_idx:end_idx], play_buttons=True)