%load_ext autoreload
+%autoreload 2
+import matplotlib.pyplot as plt
+%matplotlib inline
+import spatial_maps as sp
+import septum_mec.analysis.data_processing as dp
+import septum_mec.analysis.registration
+import expipe
+import os
+import pathlib
+import numpy as np
+import exdir
+import pandas as pd
+import optogenetics as og
+import quantities as pq
+import shutil
+from distutils.dir_util import copy_tree
+import elephant as el
+import neo
+from scipy.signal import find_peaks
+from scipy.interpolate import interp1d
+from matplotlib import mlab
+
+from tqdm import tqdm_notebook as tqdm
+from tqdm._tqdm_notebook import tqdm_notebook
+tqdm_notebook.pandas()
+data_loader = dp.Data()
+actions = data_loader.actions
+project = data_loader.project
+output = pathlib.Path('output/stimulus-spike-lfp-response')
+(output / 'figures').mkdir(parents=True, exist_ok=True)
+identify_neurons = actions['identify-neurons']
+# sessions = pd.read_csv(identify_neurons.data_path('sessions'))
+units = pd.read_csv(identify_neurons.data_path('units'))
+def get_lim(action_id):
+ stim_times = data_loader.stim_times(action_id)
+ if stim_times is None:
+ return [0, np.inf]
+ stim_times = np.array(stim_times)
+ return [stim_times.min(), stim_times.max()]
+
+def get_mask(lfp, lim):
+ return (lfp.times >= lim[0]) & (lfp.times <= lim[1])
+
+def zscore(a):
+ return (a - a.mean(0)) / a.std(0)
+
+def compute_stim_freq(action_id):
+ stim_times = data_loader.stim_times(action_id)
+ if stim_times is None:
+ return np.nan
+ stim_times = np.array(stim_times)
+ return 1 / np.mean(np.diff(stim_times))
+def signaltonoise(a, axis=0, ddof=0):
+ a = np.asanyarray(a)
+ m = a.mean(axis)
+ sd = a.std(axis=axis, ddof=ddof)
+ return np.where(sd == 0, 0, m / sd)
+
+
+def clean_lfp(anas, width=500, threshold=2):
+ anas = np.array(anas)
+
+ for ch in range(anas.shape[1]):
+ idxs, = np.where(abs(anas[:, ch]) > threshold)
+ for idx in idxs:
+ anas[idx-width:idx+width, ch] = 0 # TODO AR model prediction
+ return anas
+def compute_energy(p, f, f1, f2):
+ if np.isnan(f1) or np.all(np.isnan(p)):
+ return np.nan
+ mask = (f > f1) & (f < f2)
+ df = f[1] - f[0]
+ return np.sum(p[mask]) * df
+def find_theta_peak(p, f, f1, f2):
+ if np.all(np.isnan(p)):
+ return np.nan, np.nan
+ mask = (f > f1) & (f < f2)
+ p_m = p[mask]
+ f_m = f[mask]
+ peaks, _ = find_peaks(p_m)
+ idx = np.argmax(p_m[peaks])
+ return f_m[peaks[idx]], p_m[peaks[idx]]
+def compute_half_width(p, f, m_p, m_f):
+ if np.isnan(m_p):
+ return np.nan, np.nan
+ m_p_half = m_p / 2
+ half_p = p - m_p_half
+ idx_f = np.where(f <= m_f)[0].max()
+ idxs_p1, = np.where(np.diff(half_p[:idx_f + 1] > 0) == 1)
+ if len(idxs_p1) == 0:
+ return np.nan, np.nan
+ m1 = idxs_p1.max()
+ idxs_p2, = np.where(np.diff(half_p[idx_f:] > 0) == 1)
+ m2 = idxs_p2.min() + idx_f
+ assert p[m1] < m_p_half < p[m1+1], (p[m1], m_p_half, p[m1+1])
+ assert p[m2] > m_p_half > p[m2+1], (p[m2], m_p_half, p[m2+1])
+
+ f1 = interp1d([half_p[m1], half_p[m1 + 1]], [f[m1], f[m1 + 1]])(0)
+ f2 = interp1d([half_p[m2], half_p[m2 + 1]], [f[m2], f[m2 + 1]])(0)
+ return f1, f2
+def compute_stim_peak(p, f, s_f):
+ if np.isnan(s_f):
+ return np.nan
+ return interp1d(f, p)(s_f)
+def compute_spike_lfp(anas, sptr, t_start, t_stop, NFFT):
+
+ t_start = t_start * pq.s if t_start is not None else 0 * pq.s
+ sampling_rate = anas.sampling_rate
+ units = anas.units
+ if t_start is not None and t_stop is not None:
+ t_stop = t_stop * pq.s
+ mask = (anas.times > t_start) & (anas.times < t_stop)
+ anas = np.array(anas)[mask,:]
+
+ cleaned_anas = zscore(clean_lfp(anas))
+
+ cleaned_anas = neo.AnalogSignal(
+ signal=cleaned_anas * units, sampling_rate=sampling_rate, t_start=t_start
+ )
+
+ sptr = neo.SpikeTrain(
+ sptr.times[(sptr.times > t_start) & (sptr.times < cleaned_anas.times[-1])],
+ t_start=t_start, t_stop=cleaned_anas.times[-1]
+ )
+
+ sigs, freqs = el.sta.spike_field_coherence(cleaned_anas, sptr, **{'nperseg': NFFT})
+ return sigs, freqs, cleaned_anas
+# action_id_0 = '1833-200619-1'
+# action_id_1 = '1833-200619-2'
+# lfp_0 = data_loader.lfp(action_id_0, 6)
+# lfp_1 = data_loader.lfp(action_id_1, 6)
+
+# lim_0 = get_lim(action_id_0)
+# lim_1 = get_lim(action_id_1)
+
+# sptrs_0 = data_loader.spike_trains(action_id_0, 6)
+# sptrs_1 = data_loader.spike_trains(action_id_1, 6)
+
+# coher_0, freq_0, clean_lfp_0 = compute_spike_lfp(lfp_0, sptrs_0[163], *lim_0, 4096)
+# coher_1, freq_1, clean_lfp_1 = compute_spike_lfp(lfp_1, sptrs_1[28], *lim_1, 4096)
+
+# best_channel_0 = np.argmax(signaltonoise(clean_lfp_0))
+# best_channel_1 = np.argmax(signaltonoise(clean_lfp_1))
+
+# plt.plot(freq_0, coher_0[:,best_channel_0])
+# plt.plot(freq_1, coher_1[:,best_channel_1])
+# plt.xlim(0,20)
+NFFT = 4096*2
+theta_band_f1, theta_band_f2 = 6, 10
+def process(row):
+ action_id = row['action']
+ channel_group = row['channel_group']
+ unit_name = row['unit_name']
+
+ lfp = data_loader.lfp(action_id, channel_group)
+
+ sptr = data_loader.spike_train(action_id, channel_group, unit_name)
+
+ lim = get_lim(action_id)
+
+ p_xys, freq, clean_lfp = compute_spike_lfp(lfp, sptr, *lim, NFFT=NFFT)
+
+ snls = signaltonoise(clean_lfp)
+ best_channel = np.argmax(snls)
+ snl = snls[best_channel]
+ p_xy = p_xys[:,best_channel]
+
+ theta_f, theta_p_max = find_theta_peak(p_xy, freq, theta_band_f1, theta_band_f2)
+
+ theta_energy = compute_energy(p_xy, freq, theta_band_f1, theta_band_f2) # theta band 6 - 10 Hz
+
+ theta_half_f1, theta_half_f2 = compute_half_width(p_xy, freq, theta_p_max, theta_f)
+
+ theta_half_width = theta_half_f2 - theta_half_f1
+
+ theta_half_energy = compute_energy(p_xy, freq, theta_half_f1, theta_half_f2) # theta band 6 - 10 Hz
+
+ # stim
+
+ stim_freq = compute_stim_freq(action_id)
+
+ stim_p_max = compute_stim_peak(p_xy, freq, stim_freq)
+
+ stim_half_f1, stim_half_f2 = compute_half_width(p_xy, freq, stim_p_max, stim_freq)
+ stim_half_width = stim_half_f2 - stim_half_f1
+
+ stim_energy = compute_energy(p_xy, freq, stim_half_f1, stim_half_f2)
+
+ result = pd.Series({
+ 'signal_to_noise': snl,
+ 'best_channel': best_channel,
+ 'theta_freq': theta_f,
+ 'theta_peak': theta_p_max,
+ 'theta_energy': theta_energy,
+ 'theta_half_f1': theta_half_f1,
+ 'theta_half_f2': theta_half_f2,
+ 'theta_half_width': theta_half_width,
+ 'theta_half_energy': theta_half_energy,
+ 'stim_freq': stim_freq,
+ 'stim_p_max': stim_p_max,
+ 'stim_half_f1': stim_half_f1,
+ 'stim_half_f2': stim_half_f2,
+ 'stim_half_width': stim_half_width,
+ 'stim_energy': stim_energy
+ })
+ return result
+results = units.merge(
+ units.progress_apply(process, axis=1),
+ left_index=True, right_index=True)
+%debug
+action = project.require_action("stimulus-spike-lfp-response")
+action.modules['parameters'] = {
+ 'NFFT': NFFT,
+ 'theta_band_f1': theta_band_f1,
+ 'theta_band_f2': theta_band_f2
+}
+action.data['results'] = 'results.csv'
+results.to_csv(action.data_path('results'), index=False)
+copy_tree(output, str(action.data_path()))
+septum_mec.analysis.registration.store_notebook(action, "10-calculate-stimulus-spike-lfp-response.ipynb")
+
+