3 nb?@s~dZddlZddlZddljZddlm Z ddl Z dddZ dd d Z dd d ZdddZddZddZddZddZdS)z2 Created on Fri Apr 16 18:10:23 2021 @author: dB N)GridSpecc, s|j|jdk}dd|jD}||jd|d<|jdgjdjddd>jj}tj |d|dd |d<|jddgdd gd }t j |j |j f} |jj|jjjd d g} tj | d|dd | d<| jddgd dgd } | jj} |jj} t j| |jt} g}xR|D]J}||j|kjj}| j|}|t jd t|d|td}|j|qWt j |}t j || | f}tjd?d|jj}d@|j|jdf<dA|j|j df<|j|j}d|d<|j|djdf|j|djdf<d|j|jdkdf<d|j|jd kdf<d |d<t j |j|ddkj!d |j|ddkdf<|d|d<d|j|jdk|j"dBk@df<d|j|jdk|j"dCk@df<||ddf}|j!d }t#ddddgddgd d}|j$dd|j$ddtj%|dD}x`t&|D]T}||ddf}t j'|d kd }|j(||j)||dgt|d|d d!qWtj%|dE|d"}|j*d#x<|j+D]2}|\}}} |j,|| g|d$|d$g|dd%q|W|j-d&d'dFd)tj.d*d+tj%|dG|d,}!|!j*d#t j/tj0d-j1d.|j|jd kdfjt2j3t4fd/d0t&tD}"dddd d g}#x<|#D]4}$||j5|$kjj}%tj6dH|t j|%|$d2d3qTWxt&|D]}|j|}&|&d}'|j|}(d })|(j5dks|(j5d kr|"|(j})n4|(j5d krd4})n"|(j5dkrd5})n|(j5dkrd6})||'||&}*}+tj,d dg|*|+|)d7tj6|(j7d8d8|*|+d|(j8d9d:d;d<qW|j9d=d*d+|j:gtj;dS)Iz. Plots a spike raster for all neurons ORNcSsg|]}d|kr|qS)dur).0crr..\utils\make_output_figures.py sz$plot_sim_spikes2..Zmean_dur_odorsglomNT) categoriesorderedr) ascendinguPNmPNaltype)dtype )figsizepolaritytext_labLN text_lab_posrastypeeLNiLN) height_ratios width_ratioswspaceg)hspace|k) markersizecolor)sharexoffg?)labelalphaZodorsz upper left)\(?)titlelocbbox_to_anchorr)size)shareyzaxes.prop_cycler(c3s&|]}||tfVqdS)N)len)ri) colors_stdpn_glomsrr bsz#plot_sim_spikes2..g?right)haz0.5redblue)y1y2r(rleftcenter)r9vafontsizezTime (s))rrrCr r rC)r r )rr )r-r)r rg)<copy neur_typecolumnsmeangroupby sort_valuesindexvaluespd Categoricalnp concatenateeLNposiLNpos df_neur_idsrisinr value_countsceilastypeintr/linspacer3appendpltfigureLNposiloc reset_index duplicatedarangeshaperrupdatesubplotrangewhereplotdtaxis odor_list fill_betweenlegendxticksarrayrcParamsby_keystruniquedictrtextrr set_xlabel set_yticks tight_layout),simSpikesdf_AL_activityZ subsamplingmsizedf_orn_activityZdur_colsZgloms_in_orn_act_orderdf_ornsZln_orderdf_pnsZpn_orderZ orn_val_cntsZorn_val_cnts_subsampleZorn_indices_subsampledgZ orn_indicesZorn_n_subsamplesZ takespaceZ rorderingrRZSpikes_subsampledZN_CELLSgsax2r4Zinput_iZ where_spike_iax1row odor_name odor_startodor_endax3Z glom_colorsntypesntZpos_ntg1g2Z neur_entryr r<r=r)r5r6r plot_sim_spikes2s&      ,0     2   (    "      0 rd-C6?c Cs|jd}|jd}g}xbtt|D]R}ttd||d}ttt|||d}|j|||j||q&Wtj |||S)z Helper for get_AL_psths to compute peristimulus time histograms This function returns a smoothed spike array given a desired window length rr) sumrardr3rWmaxminrYrNrm) Z spike_arraywlrgZ spike_sumZn_neursresxr4lbubrrr get_PSTH_from_spike_arrays  "rc Cs$|jj}|j}g}g}x|D]}||jdk|j|k@jj}||jdk|j|k@jj} |jt||ddf|j |d|jt|| ddf|j |dqWt j |}t j |}t j t |j} t||jddf} t |jdkrt||jddf} t||jddf} ||| | fS)zp Helper for plot_AL_psths, returning the peristimulus time histograms for ORNs, PNs, eLNs, and uPNs rrN)rgrr)rRrD glom_namesrr rJrKrYrrgrNrmzerosr3timer\rPrQ) rwrxrrR hemi_gloms orn_psthsZ upn_psthsr Z glom_orn_posZ glom_upn_pos eln_psths iln_psthsrrr get_AL_psthss"  "& rr@cCs|j}tjdddd\}\}} } |tj|jdddd} |tj|jdddd} x`tt| D]P} | | }| |kr|nd}tj||kdd}|j |j ||ddf|d qjWx`tt| D]P} | | }| |kr|nd}tj||kdd}| j |j ||ddf|d qW| j |j |d j t|j d | j |j |d j t|j d |jd | jd | jd|jd|ddddd| jd|ddddd| jdddd|jd| jd| jd| jd|S)z? Plots peristimulus time histograms for select neurons rr r)rNrr)r+z iLNs ({})z eLNs ({})ZORNsZPNsZLNsztop z activez upper leftRQ?)r.r/r0 borderaxespad)r/r0rzfiring rate (Hz)ztime (s))rrrCrC)rr )rr )rr )rrZsubplotsrNargsortrrdr3rerfrformatrQrP set_titlerk set_ylabelrt)rpn_psthsrrrwZtopkrfigrrrZordered_orn_glomsZordered_pn_glomsr4r labZhemi_irrr plot_AL_psthss6""   rcCs0t||\}}}}t|||||dtjdS)z# Wrapper for plot_AL_psths r@N)rrrZrv)rwrxrrrrrrr plot_psthssrcCsR||d}||jj||ddf}|jd} ||jddf} ||jddf} ||jddf} ||jddf} tdddddgdd}tj|d}|j dx>|j D]4}|\}}}|j ||g|j d |j d g|dd qWtj d d tj|d|d }d}tj|j| dddddtj|j| jdd|ddtj|j| jdd|ddtj|j| jdd|ddtj|j| jdd|ddtjdtjdd$ddtj|d|d }|j|j||ddfdddtj||ddfdkd}|j|}x |D]}tj|ddd d!qWtjd"tjd#|||fS)%zI Plots current inputs for a given neuron, broken up by cell type neur_inputsNrrr g?)r!r$r*g?)r+r,r)r1)r)g?g-q=r&tot)r,r(r+r)r,r+rrPNzI (pA)z upper left)\(?)r/r0rri)r(-g?)r(lsr,ztime (s)zV (mV))rr )TrORNposrQrPPNposrrZrcrhrirjnALrlrfrylabelrkrNreaxvlinexlabel)neur_i d_neur_inputsrwIIinVrxrZneur_i_I_by_inputsZ neur_i_I_totZneur_i_inputs_ORNZneur_i_inputs_iLNZneur_i_inputs_eLNZneur_i_inputs_PNrZax_odorrrrrrZalphrZ wherespikesZ spiketimesstrrr plot_neur_I_contribss@     *  "   rcCsDddddg}ddddg}|jj}|j}tj||}||jd d fd d |jfj}||jd d fd d |jfj}||jd d fd d |jfj} ||jd d fd d |j fj} ||jd d fd d |jfj} ||jd d fd d |jfj} ||jd d fd d |jfj} ||jd d fd d |j fj}||jd d fd d |jfj}||jd d fd d |jfj}||jd d fd d |jfj}||jd d fd d |j fj}||j d d fd d |jfj}||j d d fd d |jfj}||j d d fd d |jfj}||j d d fd d |j fj}||| | | | | |||||||||g}d d |D}tj d d |Dj d d }t j |||d}||_||_d|j_d|j_tjdtjdtj|}d||dk<tj|ddddddd|d tjddd S)z Plots a 4x4 grid indicating the cell type - cell type multiplier applied to the connectivity matrix of a Sim object (i.e. what is the multiplier applied to each iLN -> PN synapse?) rrrror4epNcSs"g|]}tj|tj|qSr)rNrqisnan)rxrrr r sz+plot_synapse_scale_hmap..cSsg|]}tj|dqS)r)rNabs)rrrrr r s)rJrFfromtoequalzs   s   )1)