Showing Receptive Fields

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 .

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

%matplotlib inline

dandiset_id = "000021"
dandi_filepath = "sub-716813540/sub-716813540_ses-739448407.nwb"
download_loc = "."
dandi_api_key = None

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

units = nwb.units
units[:10]

nwb.intervals.keys()

stim_name = "gabors_presentations"
# stim_name = "RF mapping_presentations"
rf_stim_table = nwb.intervals[stim_name].to_dataframe()
rf_stim_table[:10]

### helper functions made from electrodes table to get brain location of unit or probe name of unit

# In older Allen NWBs the information must be manually cross-referenced between the units table and electrodes table

# map channel ids to brain location and probe group name
# creates a dictionary that iterates over indices of the dataset, assigns "id" as the key, and makes the location the value
channel_locations = {nwb.electrodes["id"][i]: nwb.electrodes["location"][i] for i in range(len(nwb.electrodes))}
# creates similar dictionary as previous line, but with the group name/probe name as the value
channel_probes = {nwb.electrodes["id"][i]: nwb.electrodes["group_name"][i] for i in range(len(nwb.electrodes))}

# function retrieves peak channel ID for given unit index from "units" dataset and looks up corresponding location from dictionary
def get_unit_location(unit_idx):
    peak_channel_id = units["peak_channel_id"][unit_idx]
    return channel_locations[peak_channel_id]

# function retrieves peak channel ID from "units" dataset, looks up corresponding group name, returns probe associated with peak channel
def get_unit_probe(unit_idx):
    peak_channel_id = units["peak_channel_id"][unit_idx]
    return channel_probes[peak_channel_id]



# In newer Allen NWBs, the units table includes direct references to each unit's electrodes, so we can access them directly
# def get_peak_channel_idx(unit_idx):
#     mean_waveforms = units["waveform_mean"][unit_idx]
#     waveform_mins = np.min(mean_waveforms, axis=0)
#     peak_channel_idx = np.argmin(waveform_mins)
#     return peak_channel_idx

# def get_unit_location(unit_idx):
#     peak_channel_idx = get_peak_channel_idx(unit_idx)
#     detected_electrodes = units["electrodes"][unit_idx]
#     return detected_electrodes.iloc[peak_channel_idx].location

# def get_unit_probe(unit_idx):
#     peak_channel_idx = get_peak_channel_idx(unit_idx)
#     detected_electrodes = units["electrodes"][unit_idx]
#     return detected_electrodes.iloc[peak_channel_idx].group_name

print(set(channel_locations.values()))
print(set(channel_probes.values()))

### select units to view the receptive fields of, change the conditional function below to select different units

isi_column = "isi_violations_count" if "isi_violations_count" in units.colnames else "isi_violations"

def select_condition(unit_idx):
    # the values below are recommended thresholds for these quality metrics
    return units["snr"][unit_idx] > 1 and \
            units[isi_column][unit_idx] < 1 and \
            units["firing_rate"][unit_idx] > 0.1
            # get_unit_location(unit_idx) == "VISam"
            # get_unit_probe(unit_idx) == "probeC"
            # change the line above an option above to select a different brain location or probe

selected_unit_idxs = []
for unit_idx in range(len(units)):
    if select_condition(unit_idx):
        selected_unit_idxs.append(unit_idx)
        
if len(selected_unit_idxs) == 0:
    raise IndexError("There are no units for this selection")

print(len(selected_unit_idxs))

### get x and y coordinates of gabors displayed to build receptive field

for coord_key in ('pos_x','x_position','x_pos','X'):
	if coord_key in rf_stim_table.columns:
		x_key = coord_key
		y_key = coord_key.replace('x','y').replace('X','Y')
		break
else:
	raise ValueError("Could not find x and y columns in RF intervals table")


xs = np.sort(list(set(rf_stim_table[x_key])))
ys = np.sort(list(set(rf_stim_table[y_key])))
field_units = rf_stim_table.get("units", ["unknown"])[0]
print(xs)
print(ys)
print(field_units)

def get_rf_vectorized(spike_times, rf_stim_table, xs=None, ys=None, response_window=0.2):
    """
    Calculate receptive field using vectorized operations for improved performance.
    
    Parameters:
    -----------
    spike_times : array-like
        Spike times for the unit
    rf_stim_table : pandas.DataFrame
        Stimulus table containing gabor presentations
    xs : array-like, optional
        X-coordinates of the receptive field. If not provided, uses rf_stim_table.x_positions
    ys : array-like, optional
        Y-coordinates of the receptive field. If not provided, uses rf_stim_table.y_positions
    response_window : float, optional
        Time window after stimulus onset to count spikes (default: 0.2 seconds)
        
    Returns:
    --------
    unit_rf : numpy.ndarray
        2D array containing spike counts for each position in the receptive field
    """
    if xs is None:
        xs = np.sort(rf_stim_table[x_key].unique())
    if ys is None:
        ys = np.sort(rf_stim_table[y_key].unique())
    
    # Initialize the receptive field array
    unit_rf = np.full((len(ys), len(xs)), np.nan)
    
    # Pre-compute stimulus presentations by position
    position_stim_times = {}
    for xi, x in enumerate(xs):
        for yi, y in enumerate(ys):
            position_key = (x, y)
            position_stim_times[position_key] = rf_stim_table[
                (rf_stim_table[x_key] == x) & 
                (rf_stim_table[y_key] == y)
            ].start_time.values
    
    # Process each position
    for xi, x in enumerate(xs):
        for yi, y in enumerate(ys):
            position_key = (x, y)
            stim_times = position_stim_times[position_key]
            
            if len(stim_times) == 0:
                continue
                
            # Create arrays of start and end times for the response windows
            start_times = stim_times
            end_times = stim_times + response_window
            
            # Find indices for all spikes within the response windows in one operation
            # This is much faster than searching for each stimulus individually
            all_indices = np.searchsorted(spike_times, np.column_stack((start_times, end_times)))
            
            # Calculate spike counts for all presentations at once
            spike_counts = all_indices[:, 1] - all_indices[:, 0]
            
            # Store the total spike count for this position
            unit_rf[yi, xi] = np.mean(spike_counts)
    
    return unit_rf

### compute receptive fields for each unit in selected units

unit_rfs = []
for unit_idx in selected_unit_idxs:
    unit_spike_times = units["spike_times"][unit_idx]
    unit_rfs.append(get_rf_vectorized(unit_spike_times, rf_stim_table))

### display the receptive fields for each unit in a 2D plot

n_rows = len(unit_rfs)//10
fig, axes = plt.subplots(n_rows+1, 10)
fig.set_size_inches(12, n_rows+1)

# handle case where there's <= 10 rfs
if len(axes.shape) == 1:
    axes = axes.reshape((1, axes.shape[0]))

for irf, rf in enumerate(unit_rfs):
    ax_row = int(irf/10)
    ax_col = irf%10
    axes[ax_row][ax_col].imshow(rf, origin="lower")
for ax in axes.flat[1:]:
    ax.axis('off')

# making axis labels for first receptive field
axes[0][0].set_xlabel(field_units)
axes[0][0].set_ylabel(field_units)
axes[0][0].xaxis.set_label_position("top") 
axes[0][0].xaxis.set_ticks_position("top")
axes[0][0].set_xticks(range(len(xs)), xs, rotation=90, fontsize=6)
axes[0][0].set_yticks(range(len(ys)), ys, fontsize=6)

for i, l in enumerate(axes[0][0].xaxis.get_ticklabels()):
    if i % 2 != 0:
        l.set_visible(False)
for i, l in enumerate(axes[0][0].yaxis.get_ticklabels()):
    if i % 2 != 0:
        l.set_visible(False)