Download and plot power spectrum of raw data

Downloads LFP power spectrum for a given session and probe and plots a heatmap of power spectrum on the channels along probe against frequency

# import modules
from one.api import ONE
import matplotlib.pyplot as plt
import numpy as np

# instantiate ONE
one = ONE(base_url='https://openalyx.internationalbrainlab.org', silent=True)

# Specify subject, date and probe we are interested in
subject = 'CSHL049'
date = '2020-01-08'
sess_no = 1
probe_label = 'probe00'
eid = one.search(subject=subject, date=date, number=sess_no)[0]

# Download the data
# channels.rawInd: Index of good recording channels along probe
# channels.localCoordinates: Position of each recording channel along probe
channels = one.load_object(eid, 'channels', collection=f'alf/{probe_label}')
# Get range for y-axis
depth_range = [np.min(channels.localCoordinates[:, 1]),
               np.max(channels.localCoordinates[:, 1])]

# Load in power spectrum data
lfp_spectrum = one.load_object(eid, 'ephysSpectralDensityLF',
                               collection=f'raw_ephys_data/{probe_label}')
lfp_freq = lfp_spectrum['freqs']
lfp_power = lfp_spectrum['power'][:, channels.rawInd]

# Define a frequency range of interest
freq_range = [0, 300]
freq_idx = np.where((lfp_freq >= freq_range[0]) &
                    (lfp_freq < freq_range[1]))[0]

# Limit data to freq range of interest and also convert to dB
lfp_spectrum_data = 10 * np.log(lfp_power[freq_idx, :])
dB_levels = np.quantile(lfp_spectrum_data, [0.1, 0.9])

# Create figure
fig, ax = plt.subplots()
# Plot the LFP spectral data
spectrum_plot = ax.imshow(lfp_spectrum_data.T, extent=np.r_[freq_range, depth_range],
                          cmap='viridis', vmin=dB_levels[0], vmax=dB_levels[1], origin='lower',
                          aspect='auto')
cbar = fig.colorbar(spectrum_plot, ax=ax)
cbar.set_label('LFP power (dB)')
ax.set_xlabel('Frequency (Hz)')
ax.set_ylabel('Depth along probe (um)')
ax.set_title('Power Spectrum of LFP')

plt.show()