<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> from phy . utils . _types import _as_array <EOL> from phy . io . array import _index_of , _unique <EOL> def _increment ( arr , indices ) : <EOL> """<STR_LIT>""" <EOL> arr = _as_array ( arr ) <EOL> indices = _as_array ( indices ) <EOL> bbins = np . bincount ( indices ) <EOL> arr [ : len ( bbins ) ] += bbins <EOL> return arr <EOL> def _diff_shifted ( arr , steps = <NUM_LIT:1> ) : <EOL> arr = _as_array ( arr ) <EOL> return arr [ steps : ] - arr [ : len ( arr ) - steps ] <EOL> def _create_correlograms_array ( n_clusters , winsize_bins ) : <EOL> return np . zeros ( ( n_clusters , n_clusters , winsize_bins // <NUM_LIT:2> + <NUM_LIT:1> ) , <EOL> dtype = np . int32 ) <EOL> def _symmetrize_correlograms ( correlograms ) : <EOL> """<STR_LIT>""" <EOL> n_clusters , _ , n_bins = correlograms . shape <EOL> assert n_clusters == _ <EOL> correlograms [ ... , <NUM_LIT:0> ] = np . maximum ( correlograms [ ... , <NUM_LIT:0> ] , <EOL> correlograms [ ... , <NUM_LIT:0> ] . T ) <EOL> sym = correlograms [ ... , <NUM_LIT:1> : ] [ ... , : : - <NUM_LIT:1> ] <EOL> sym = np . transpose ( sym , ( <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:2> ) ) <EOL> return np . dstack ( ( sym , correlograms ) ) <EOL> def correlograms ( spike_times , <EOL> spike_clusters , <EOL> cluster_ids = None , <EOL> sample_rate = <NUM_LIT:1.> , <EOL> bin_size = None , <EOL> window_size = None , <EOL> symmetrize = True , <EOL> ) : <EOL> """<STR_LIT>""" <EOL> assert sample_rate > <NUM_LIT:0.> <EOL> assert np . all ( np . diff ( spike_times ) >= <NUM_LIT:0> ) , ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> spike_times = np . asarray ( spike_times , dtype = np . float64 ) <EOL> spike_samples = ( spike_times * sample_rate ) . astype ( np . int64 ) <EOL> spike_clusters = _as_array ( spike_clusters ) <EOL> assert spike_samples . ndim == <NUM_LIT:1> <EOL> assert spike_samples . shape == spike_clusters . shape <EOL> bin_size = np . clip ( bin_size , <NUM_LIT> , <NUM_LIT> ) <EOL> binsize = int ( sample_rate * bin_size ) <EOL> assert binsize >= <NUM_LIT:1> <EOL> window_size = np . clip ( window_size , <NUM_LIT> , <NUM_LIT> ) <EOL> winsize_bins = <NUM_LIT:2> * int ( <NUM_LIT> * window_size / bin_size ) + <NUM_LIT:1> <EOL> assert winsize_bins >= <NUM_LIT:1> <EOL> assert winsize_bins % <NUM_LIT:2> == <NUM_LIT:1> <EOL> if cluster_ids is None : <EOL> clusters = _unique ( spike_clusters ) <EOL> else : <EOL> clusters = _as_array ( cluster_ids ) <EOL> n_clusters = len ( clusters ) <EOL> spike_clusters_i = _index_of ( spike_clusters , clusters ) <EOL> shift = <NUM_LIT:1> <EOL> mask = np . ones_like ( spike_samples , dtype = np . bool ) <EOL> correlograms = _create_correlograms_array ( n_clusters , winsize_bins ) <EOL> while mask [ : - shift ] . any ( ) : <EOL> spike_diff = _diff_shifted ( spike_samples , shift ) <EOL> spike_diff_b = spike_diff // binsize <EOL> mask [ : - shift ] [ spike_diff_b > ( winsize_bins // <NUM_LIT:2> ) ] = False <EOL> m = mask [ : - shift ] . copy ( ) <EOL> d = spike_diff_b [ m ] <EOL> d = spike_diff_b [ m ] <EOL> indices = np . ravel_multi_index ( ( spike_clusters_i [ : - shift ] [ m ] , <EOL> spike_clusters_i [ + shift : ] [ m ] , <EOL> d ) , <EOL> correlograms . shape ) <EOL> _increment ( correlograms . ravel ( ) , indices ) <EOL> shift += <NUM_LIT:1> <EOL> correlograms [ np . arange ( n_clusters ) , <EOL> np . arange ( n_clusters ) , <EOL> <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> if symmetrize : <EOL> return _symmetrize_correlograms ( correlograms ) <EOL> else : <EOL> return correlograms </s>
