a
    Fgb                     @   s   d Z ddlZddlZddlmZ ddlZddlm	Z	 ddl
ZddlZdd Zd ddZd	d
 Zd!ddZdd Zd"ddZd#ddZdd Zdd Zdd Zedkre Ze  ede e   e  dS )$zn
Created on Wed Apr 24 23:05:58 2024

@author: Albert

Functions to perform TCA, LDA, and one after the other
    N)LinearDiscriminantAnalysisc                   C   s
   t   d S )N)test_class_inequality r   r   QD:\MPI_Brain\Codes\aversive_scripts_2025\aversive_scripts_2025\APdecoding_funs.pymain   s    r   Fc              
   C   s   d}||k rz,t j|d}	|	j| td|d ||d W q tyR   |d7 }Y q tjjyn   |d7 }Y q ty }
 zt	d|
 |
W Y d}
~
qd}
~
0 0 qq||krt	d |	
|| }|s|S ||	fS dS )a$   Performs TCA on a data matrix 
        data_by_trial: array of size "num features" X "num time bins" X "num trials"
        TCA_factors: number of TCA dimensions for each factor
        TCA_replicates: number of ensembles (recommended high, otherwise it won't converge well)
        TCA_method: "cp_als", "mcp_als", "ncp_bcd", "ncp_hals"
        TCA_convergence_attempts: how many times to try TCA. Will raise an error if it hasn't converged after enough attempts
        TCA_replicate_selection: index of the replicate to get (defaults to 0)
    r   )
fit_method   )ranks
replicatesverbosezUnexpected errorNzWARNING: TCA did not converge!)tensortoolsEnsemblefitrange
IndexErrornplinalgLinAlgError	Exceptionprintfactors)data_by_trialTCA_factorsTCA_replicates
TCA_methodTCA_convergence_attemptsTCA_replicate_selectionr   return_ensembleTCA_attempts_counterTCA_ensembleerrorKTensorr   r   r   perform_TCA   s&    
 
r"   c                 C   s~   t |}t |}t|D ]*\}}|d|}| | }|||< q|t j| }t 	d| | d|  szJ |S )Nr      )
r   eyezeros	enumeratereshape	transformsqueezer   normallclose)LDA_objnum_featuresIDLDA_axdim_idxeuler_vector
euler_projr   r   r   get_LDA_axisE   s    


"r3   r   333333?   c           :         s  | j d }t t } fdd|D fddtt|D }t|t| }}	i }
t|f}tjtd}t}|	|kstj	t
dg}d}ng }tttt }t|t||  }tt|| D ]}|| } |k} |k}t}t|d }tjj|t|dd}t|D ]}||| |d |  }t||kr|d	||  }tjj||t| dd}t||f}t|}d
||< || qDqtj	|ftd }t||f}t|f}t|t|f}t|t|f} t||fd }!t|D ]}|| }"t|"}#t|#|f}$t|#D ]}%t|"}&t|"d |% }'d|&|'< | |& }( |& })| |'g }* |' }+tdd|dd},|,|(|) t|,|*|||'f< |,|*|||'f< |,|*d|+f |||'f< t|,j|$|%d	d	f< qf |" }-|||"f }.t|-|.}/|/||< g }0|D ]8}1|-|1k}2t|.|2 |-|2 kt|2 }3|0|3 q\|0| |< tj|!| dd|!|< tj|$dd|!|< q8tj|dd}4tj| dd}5tj|!dd}6tD ]r}%|d	d	|%f fdd|D }7|t|7 }8|8k}9t||9 ||%d	d	f< t||9 ||%< |8||%< q|4|
d< |5|
d< ||
d< ||
d< |6|
d< ||
d< |
S )aT   LDA_input: size "num samples" X "num features"
        labels: size "num samples". ASSUMES THEY ARE "INT" THAT START AT 0 AND INCREASE 1 BY 1, NO NEGATIVES OR DECIMALS
        LDA_components: int
        imbalance_prop: float between 0 and 1 (if the proportion of a class is larger than this, deal with data imbalance during training)
    r   c                    s   g | ]}t  |kqS r   r   sum.0label)labelsr   r   
<listcomp>\       zperform_LDA.<locals>.<listcomp>c                    s   g | ]}|   qS r   r   )r9   lidx)num_samplessamples_by_labelr   r   r<   ]   r=   dtyper   F)replaceNTeigensolver	shrinkagen_componentsstore_covariance)axisc                    s   g | ]}t  |kqS r   r6   r8   )sample_labelsr   r   r<      r=   f1accuracylabel_predictedLDA_probweightsLDA_projection)shapelenr   uniquer   argmaxmaxr%   intonesboolceilminwhererandomchoicehstackcopyappendr7   r   r   r)   r(   predictpredict_probaabscoef_pfmulticlass_f1average):	LDA_inputr;   LDA_componentsimbalance_propimbalance_repetitionsr-   unique_labelsprop_by_labelmajority_class_idxmajority_propLDA_results_dictrQ   rN   LDA_prob_correctselected_trials_list,min_selections_needed_to_include_all_samplesrepmajority_classmajority_boolminority_boolminority_samplesmajority_idxsmajority_idxs_shuffledsubsetmajority_idxs_selectedalready_used_idxsfiller_idxsselected_trialslabel_by_repproj_by_repprob_by_rep	f1_by_rep
acc_by_repweights_by_repselected_trials_boolnum_samples_by_repweights_by_fold
sample_idxselected_trials_bool_currenttest_idxLDA_train_inputLDA_train_labelLDA_test_inputLDA_test_labelLDAselected_labelspred_labelsrL   acc_listr:   
label_boolaccf1_avgacc_avgrP   prediction_label_countmost_predicted_label rep_idxs_of_most_predicted_labelr   )r;   r?   rK   r@   r   perform_LDAQ   s    







 r   c                 C   s   t dd | D }|dkr:t jt|td}t| }n|dkr`t dd |D }|| }nx|dkrt dd |D }|| }d	||d
k< nF|dkrt dd |D }|| }d||d	k< d	||d
k< ntd ||fS )a   Given a list of session numbers, select and label them depending on "session_comparisons"
        snum_by_trial: list of session numbers
        session_comparison: string, can be
            - airpuff: takes all sessions, splits airpuff from non-airpuff (B and P vs T)
            - BT: compares B vs T sessions
            - BP: compares B vs P sessions
            - TP: compares T vs P sessions
    c                 S   s   g | ]}t jt j| qS r   )pparamSESSION_TYPE_LABELSindexSESSION_LABEL_BY_SNUMr9   snumr   r   r   r<      r=   z1get_trials_to_keep_and_labels.<locals>.<listcomp>airpuffrA   BTc                 S   s   g | ]}|d v qS )r   r   r   r9   styper   r   r   r<      r=   BPc                 S   s   g | ]}|d v qS ))r   r#   r   r   r   r   r   r<      r=   r   r#   TPc                 S   s   g | ]}|d v qS ))r   r#   r   r   r   r   r   r<      r=   r   z*WARNING: WRONG "session_comparisons" label)r   arrayrX   rS   rY   r    get_AP_labels_from_snum_by_trialr   )snum_by_trialsession_comparisonsstype_by_trialtrials_to_keeplabel_by_trialr   r   r   get_trials_to_keep_and_labels   s$    	
r   Tc                 C   sl  d}|du r0t  jd }t j|dd t  }| jdkr@d}n
| jd }t|}g d}	|dkr<ddg}
|durd	d
 |D }t	j
d |v rddg}
tt|}t|D ]\}}t||kd }|| }| | }|j||d|	| |
| d |jd|d |jd|d | d dkr|| d dg qn|dkrg d}
| dddf | dddf  }}t|D ]d\}}t||kd }|| }|| }|j||d|	| |
| d |jd|d |jd|d qx|jd|d |jd|d |r|j|ddd |dur$t|| |jddg d d D ]}|j| d! q<t  }|  ||fS )"aD   Plot LDA projection. Assumes it's either 1d or 2d
        LDA_projection is the output of scipy's LDA.transform, has shape "num_trials" X "LDA dimension"
        label_by_trial has shape "num_trials"
        snum_by_trial has shape "num_trials", indicates the session number. Used to mark session separation (optional)
       Nr   )      figsize	royalblue	indianredforestgreenzLabel 1zLabel 2c                 S   s   g | ]}t j| qS r   )r   r   r   r   r   r   r<   6  r=   z'plot_LDA_projection.<locals>.<listcomp>zNo APAPr   d   )scolorr:   zTrial numberfontsizezLDA projection (AU)g      @r#   )BTPzLDA 1zLDA 2y)rJ   	labelsizexzlower rightF)r   locframeonrighttop)r   bottomleftr   r   )pltgcfnumberfiguregcandimrR   r   rT   r   r   arangerS   r&   r\   scatter
set_xlabel
set_ylabelget_ylimset_ylimtick_paramslegend$add_session_delimiters_to_trial_plotspinesset_visibleset_linewidthtight_layout)trial_factors_LDA_projectionr   r   plot_legendaxfsfig_numrj   rm   label_colorslabel_namesr   ldat	label_idxr:   
label_idxstlda1ldaxldayr   r   rJ   figr   r   r   plot_LDA_projection  s^    




"

r   r   c                 C   sh  |   \}}|d }tt|}d}d}	t|D ]\}
}tjtj| }||k}t	|}t
|t	|ddd  d }|
dkr|}	|
t
|d kr||
d  }tjtj| }nd}||kr.|	}|}	|dkr||	| d  }| j||dtj|  |d	d
 |dkr.| j||	g||ddd | j|gd tj||dddddd q4| ||g dS )z For a plot whose x axis is the trial number, add session delimiters.
        snum_by_trial: array of size "num trials", each element is the session number
        assumes sessions 3,4,5,6 have an airpuff, and are marked differently
    g333333?r   Nr   Tr#   z%sbold)r   weightgray皙?)r   alphar   )numz--      ?)r   r   sortrT   r&   r   r   r   r   rU   rS   textSESSION_TYPE_LABELS_SHORTfill_betweenplotlinspacer   )r   r   r   add_stype_labelsyminymax_originalymaxsnum_uniquestype_start_trialstype_end_trialsnum_idxr   	stype_idxsession_trial_idxssession_first_trial_idxsession_last_trial_idx	snum_nextstype_idx_next	mid_trialr   r   r   r   f  s4    
 


(r   c                  C   s.  d} d}d}d}d}d}t ||f}t dd||dd d f< t dt j t | | |dd d f< t|  | d7 } t|dd d f  t|  | d7 } t|dd d f  t ||f}t 	t jt | | |dd d f< d|dt
|d d f< t t jt | | |dd d f< d|dt
|d d f< t|  | d7 } t|dd d f  t|  | d7 } t|dd d f  t ||f}g d|dd d f< g d	|dd d f< t|  | d7 } t|dd d f d
 t|  | d7 } t|dd d f d
 t |||f}	t|D ]x}
t|D ]h}t|D ]X}t |d d |
f |d d |f  |d d |f  t jjd|dd |	|
||f< qNqBq6t|	|dddddd}|d j}t|j t|  | d7 } t|dd d f  t|  | d7 } t|dd d f  d S )Nr   
         r#   皙?r   )r4   g?ffffff?r  r   r  )g?g?r   r  g333333?g?or   scalesizencp_halsF)r   r   )r   rX   r   cospir   r   r   r   sinrW   r%   r   r7   r]   normalr"   r   r   rR   )r   
num_trials	num_timesnum_neuronsnum_factorsstdtrial_factorstime_factorsneuron_factorsdatanr   kr!   trial_factors_tcar   r   r   test_TCA_in_false_data  sN    *>>^

r'  c                  C   s  d} g d}d}d}d}d}t ||f}d|d< d|d< d|d	< d
|d< t jt d|d |ddt}t| t ||f}t|D ]*}	||	 }
||
 t jj	d||d ||	< qt
dd|d dd}||| ||}t ||k| }|jd }t| t||}t| t| t|j t|jj tj| dd | d7 } t|D ]:}||k}|| }tj||df ||df d||d qRt }|d |jd| dd ddg}tdD ]}tj| dd | d7 } t|D ]0}||k}|| }tj|||f ||dd qt }|  |dkrx|jd||  d d |jd tj|j dd! |jddd" nN|dkr|j d||  d d |j!d tj|j dd! |jddd# qd$S )%z How to extract the importance of each underlying dimension for a classification problem? 
        We put all important information in one dimension
    
    r   r   r#   r      g      ?)r   r   )r   r   r   r   )r   r   r   )decimals        r  rD   FrE   )r   r   r   r  )r   r:   offzAccuracy: %.2fr  r   Z   )r  r#   r   )r   r:   r   zLDA weight: %.2f   )visible)r   labelbottom)r   	labelleftN)"r   r%   aroundr   astyperW   r   r   r]   r  r   r   rb   r7   re   r3   rR   r   r   r   r   rJ   	set_titlehistinvert_xaxisr   yaxisr   setpr   valuesr   r   xaxis)r   colorsnum_classesr-   r?   r  centersr;   r#  sampler:   r   rN   rM   rP   r/   class_label
class_idxsr   r   	rotationsrJ   r   r   r   test_LDA_weight_function  sv    " 



(


rA  c            	   	   C   s  d} d}d}t dg|  dg|  }t jj|t|dd}d}t ddg}t|D ]v}t jj|t|dd}|t||7 }|d  t |d	|  dk|  7  < |d  t || | dk| 7  < qV|| }|| }t	| | |  || |   t	| t	| d	S )
z Test for the expected accuracy and f1 scores when classes are imbalanced 
        We compare the original imbalanced labels with completely random labels    
    r   r  i  r   r   F)r  rC   r*  N)
r   r   r]   r^   rS   r   rf   rg   r7   r   )	s1s2Nrefref_shuffler   piranr   r   r   r   5  s    &(r   __main__zTime Ellapsed: %.1f)r   FF)r   r4   r5   )NTN)r   T)__doc__timenumpyr   matplotlib.pyplotpyplotr   r   sklearn.discriminant_analysisr   project_parametersr   processing_functionsrf   r   r"   r3   r   r   r   r   r'  rA  r   __name__ttr   showr   r   r   r   <module>   s2     
'
 +
Z
/?a