# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# + [markdown] id="view-in-github" colab_type="text"
# [View in Colaboratory](https://colab.research.google.com/github/mqclark/home/blob/master/flight_analysis_script.ipynb)

# + id="jypNilpZ0PqQ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 469} outputId="6356fb0c-44da-421f-90a3-460adf929247"
# #!pip install neo==0.4

import numpy as np
import matplotlib.pyplot as plt

flydir = 'C:/Users/rancher/Desktop/abf_csv_IO/' #change directory as needed

from neo.io import AxonIO

localfile = flydir + 'test.abf' #change file name as needed

r = AxonIO(localfile)

bl = r.read_block(lazy=False, cascade=True)

l_amp = np.asarray(bl.segments[0].analogsignals[3])*60/np.pi
r_amp = np.asarray(bl.segments[0].analogsignals[4])*60/np.pi
wbf = np.asarray(bl.segments[0].analogsignals[0])*100.
#pattern = np.asarray(bl.segments[0].analogsignals[6])
#frames = np.asarray(bl.segments[0].analogsignals[3])
#add Chr_stim = XXXXXXXXX

amp_sum = l_amp + r_amp
l_amp = l_amp - r_amp

fs_axon = 1.0/2000.0
times=np.linspace(0,len(l_amp)/(1/fs_axon),len(l_amp))

ax = plt.subplot(211)
plt.plot(times, pattern)

ax2 = plt.subplot(212, sharex=ax)
ax2.plot(times, l_amp)

#ax = plt.subplot2grid((11,6),(2,0),colspan = 3)
#plt.plot(times, wbf)

thresh = 0.1

all_trials = pattern
isInTrialInds = (all_trials > thresh)

trialStartInds = (np.diff(isInTrialInds) == 1)

#find the start and end times of each stimulus epoch
trialStartTimes = times[trialStartInds]
trialbegin =  trialStartTimes[0:len(trialStartTimes):2]

trialDur = 3/fs_axon

pres_types = np.histogram(pattern,8)

anCodes = pres_types[1][1:]

for i in trialStartTimes:
    codeInd = np.abs(np.mean(pattern[(times > i) & (times < i + trialDur)]) - anCodes)


# Find the codes for each trial

TrialCodeidx = []
trialval = []
for j in range(len(trialStartTimes) - 1):
    TrialCodeidx = np.argmin(np.abs(np.mean(pattern[trialStartTimes[j]/fs_axon:trialStartTimes[j+1]/fs_axon]) - codeInd))
    trialval.append(TrialCodeidx)

trialval = trialval[0:len(trialval):2]

# find the indices that correspond to a given trial
def find(lst, a):
    return [i for i, x in enumerate(lst) if x == a]

grouped_trials = []
trialval2 = np.arange(0,8)
#trialval2.sort()
 
for j in range(8):
    trials = find(trialval, trialval2[j])
    grouped_trials.append(trials)
    
timePrev = 1/fs_axon

def sig_snip(sig, start, samp):
    #return l_amp[times[trialbegin[start]]/fs_axon - timePrev : times[trialbegin[start]]/fs_axon + trialDur + timePrev]
    return sig[trialbegin[start] / samp - 0.5 / samp : trialbegin[start] / samp + 4. / samp] 

triggered_ramp = []
triggered_lamp = []
triggered_wbf = []
triggered_sum = []

for j in range(8):    
    for i in range(5):
        r_snip = sig_snip(r_amp, grouped_trials[j][i], fs_axon)
        triggered_ramp.append(r_snip)
        l_snip = sig_snip(l_amp, grouped_trials[j][i], fs_axon)
        triggered_lamp.append(l_snip)
        w_snip = sig_snip(wbf, grouped_trials[j][i], fs_axon)
        triggered_wbf.append(w_snip)
        d_snip = sig_snip(amp_sum, grouped_trials[j][i], fs_axon)
        triggered_sum.append(d_snip)

wbf_means = []  
wbf_errors = []      
lamp_means = []
lamp_errors = []
sum_means = []
sum_errors = []
for i in range(0, 40, 5):   
    mean = np.mean(triggered_lamp[i:i+5], 0)
    lamp_means.append(mean)
    std = np.std(triggered_lamp[i:i+5], 0)
    lamp_errors.append(std)
    mean = np.mean(triggered_sum[i:i+5], 0)
    sum_means.append(mean)
    std = np.std(triggered_sum[i:i+5], 0)
    sum_errors.append(std)
    mean = np.mean(triggered_wbf[i:i+5], 0)
    wbf_means.append(mean)
    std = np.std(triggered_wbf[i:i+5], 0)
    wbf_errors.append(std)
    
    
#ramp_means = []
#ramp_errors = []
#diff_means = []
#diff_errors = []
#for i in range(0, 40, 5):   
#    mean = np.mean(triggered_ramp[i:i+5], 0)
#    ramp_means.append(mean)
#    std = np.std(triggered_ramp[i:i+5], 0)
#    ramp_errors.append(std)
#    mean = np.mean(triggered_diff[i:i+5], 0)
#    diff_means.append(mean)
#    std = np.std(triggered_diff[i:i+5], 0)
#    diff_errors.append(std)
    
#from scipy import signal
#indiv_xcorr = signal.fftconvolve(l_amp, wbf[::-1], mode = 'full')
#indiv_xcorr = indiv_xcorr/ np.max(np.abs(indiv_xcorr))
#lags = np.arange(-len(wbf)+1,len(wbf)) * fs_axon
#
#plt.figure(3)
#plt.plot(lags, indiv_xcorr, color = 'k')

#%% now plot the results for each axis
    
plot_time = np.arange(len(lamp_means[1]))
plot_time = plot_time * fs_axon
plt.figure(2, figsize = (20, 8))
#plt.subplot2grid((6,1),(0,0), aspect = 2.5)
plt.subplot(141)

plt.plot(plot_time, lamp_means[0]-np.mean(lamp_means[0][0:1000]), color = 'r')
plt.fill_between(plot_time, lamp_means[0]-np.mean(lamp_means[0][0:1000]) - lamp_errors[0], lamp_means[0] -np.mean(lamp_means[0][0:1000])+ lamp_errors[0], color = 'r', alpha = 0.3, edgecolor = 'none')
plt.plot(plot_time, lamp_means[1]-np.mean(lamp_means[1][0:1000]), color = 'b')
plt.fill_between(plot_time, lamp_means[1]-np.mean(lamp_means[1][0:1000]) - lamp_errors[1], lamp_means[1] -np.mean(lamp_means[1][0:1000])+ lamp_errors[1], color = 'b', alpha = 0.3, edgecolor = 'none')
#plt.gca().set_xticks([])
plt.ylabel('L WBA (V)')
plt.xlabel('Time (s)')
plt.title('yaw')
#plt.ylim((-2.0, 2.5))
#plt.xlim((0.0, 6.0))
plt.axvspan(0.5, 3.5, facecolor = 'gray', edgecolor = 'none', alpha = 0.3)

plt.subplot(142)
plt.plot(plot_time, lamp_means[2]-np.mean(lamp_means[2][0:1000]), color = 'r')
plt.fill_between(plot_time, lamp_means[2] -np.mean(lamp_means[2][0:1000])- lamp_errors[1], lamp_means[2]-np.mean(lamp_means[2][0:1000])+ lamp_errors[2], color = 'r', alpha = 0.3, edgecolor = 'none')
plt.plot(plot_time, lamp_means[3]-np.mean(lamp_means[3][0:1000]), color = 'b')
plt.fill_between(plot_time, lamp_means[3]-np.mean(lamp_means[3][0:1000])- lamp_errors[3], lamp_means[3]-np.mean(lamp_means[3][0:1000]) + lamp_errors[3], color = 'b', alpha = 0.3, edgecolor = 'none')
#plt.gca().set_xticks([])
plt.ylabel('R WBA (V)')
plt.title('looming')
#plt.xlim((0.0, 6.0))
plt.axvspan(0.5, 3.5, facecolor = 'gray', edgecolor = 'none', alpha = 0.3)

plt.subplot(143)

plt.plot(plot_time, lamp_means[4]-np.mean(lamp_means[4][0:1000]), color = 'r')
plt.fill_between(plot_time, lamp_means[4]-np.mean(lamp_means[4][0:1000]) - lamp_errors[4], lamp_means[4]-np.mean(lamp_means[4][0:1000]) + lamp_errors[4], color = 'r', alpha = 0.3, edgecolor = 'none')
plt.plot(plot_time, lamp_means[5]-np.mean(lamp_means[5][0:1000]), color = 'b')
plt.fill_between(plot_time, lamp_means[5]-np.mean(lamp_means[5][0:1000]) - lamp_errors[5], lamp_means[5]-np.mean(lamp_means[5][0:1000]) + lamp_errors[5], color = 'b', alpha = 0.3, edgecolor = 'none')
#plt.gca().set_xticks([])
plt.ylabel('R WBA (V)')
#plt.ylim((-2.5, 2.5))
#plt.xlim((0.0, 6.0))
plt.title('roll')
plt.axvspan(0.5, 3.5, facecolor = 'gray', edgecolor = 'none', alpha = 0.3)

plt.subplot(144)

l_amp = l_amp + r_amp
plt.plot(plot_time, sum_means[6]-np.mean(sum_means[6][0:1000]), color = 'r')
plt.fill_between(plot_time, sum_means[6]-np.mean(sum_means[6][0:1000]) - sum_errors[6], sum_means[6]-np.mean(sum_means[6][0:1000]) + sum_errors[6], color = 'r', alpha = 0.3, edgecolor = 'none')
plt.plot(plot_time, sum_means[7]-np.mean(sum_means[7][0:1000]), color = 'b')
plt.fill_between(plot_time, sum_means[7]-np.mean(sum_means[7][0:1000]) - sum_errors[7], sum_means[7]-np.mean(sum_means[7][0:1000]) + sum_errors[7], color = 'b', alpha = 0.3, edgecolor = 'none')
#plt.gca().set_xticks([])
plt.ylabel('R WBA (V)')
#plt.ylim((-2.5, 2.5))
#plt.xlim((0.0, 6.0))
plt.title('pitch')
plt.axvspan(0.5, 3.5, facecolor = 'gray', edgecolor = 'none', alpha = 0.3)


#%% now save the data

import pandas as pd

lamp_means = pd.DataFrame(np.transpose(lamp_means))
lamp_means.to_csv(flydir + 't1/lamp_means.csv', index=False)

sum_means = pd.DataFrame(np.transpose(sum_means))
sum_means.to_csv(flydir + 't1/sum_means.csv', index=False)

wbf_means = pd.DataFrame(np.transpose(wbf_means))
wbf_means.to_csv(flydir + 't1/wbf_means.csv', index=False)

# + id="PpUatOx-0PqX" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered