def detect_spikey(local_chunk,thr=None,plotose=0): if (thr is None): thr=np.repeat(50,local_chunk.shape[0]) from scipy.interpolate import interp1d nchan=local_chunk.shape[1] all_peaktimes=[]; for ich in range(nchan): all_peaktimes.extend(algo.detect_peaks(np.abs(local_chunk[:, ich]), thr[ich], valley=False, mpd=10,show=False)) def histc(X, bins): map_to_bins = np.digitize(X,bins) r = np.zeros(bins.shape) for iii in map_to_bins: r[iii-1] += 1 return [r, map_to_bins] # buf=np.random.gamma(1,3,[1000,64]) #test data all_peaktimes=np.sort(all_peaktimes) if (len(all_peaktimes)<6)|(local_chunk.shape[0]<90001): return [],local_chunk R,_=histc(all_peaktimes,np.arange(min(all_peaktimes),max(all_peaktimes),30000)) is_spikey=R> (nchan*35) # 20 spikes per second on every channel is too much if (len(is_spikey)<3): return [],local_chunk if plotose: _ = algo.detect_peaks(np.abs(local_chunk[:, -2]), thr, valley=False, mpd=10,show=True) plt.show() plt.plot(all_peaktimes,'.');plt.title('all chan over time');plt.show();plt.close() plt.plot(np.diff(all_peaktimes),'.');plt.title('all chan diff');plt.show();plt.close() #% plt.plot(R,'.');plt.title('Raster all chan over time');plt.show();plt.close() plt.plot(R/nchan,'-');plt.title('all chan spikes per chan per sec');plt.show();plt.close() plt.plot(is_spikey,'.');plt.title('is_spiky');plt.show();plt.close() f=interp1d(range(len(is_spikey)),is_spikey) is_spikey1=np.round(f(np.linspace(0,len(is_spikey)-1,local_chunk.shape[0]))) local_chunk[is_spikey1==1]=0#np.nan return is_spikey1,local_chunk
def main(params, nb_cpu, nb_gpu, use_gpu): # Part 1: Whitening numpy.random.seed(420) logger = init_logging(params.logfile) logger = logging.getLogger('circus.whitening') ################################################################# data_file = params.data_file data_file.open() N_e = params.getint('data', 'N_e') N_total = params.nb_channels N_t = params.getint('detection', 'N_t') dist_peaks = params.getint('detection', 'dist_peaks') template_shift = params.getint('detection', 'template_shift') file_out_suff = params.get('data', 'file_out_suff') file_out = params.get('data', 'file_out') spike_thresh = params.getfloat('detection', 'spike_thresh') matched_filter = params.getboolean('detection', 'matched-filter') matched_thresh = params.getfloat('detection', 'matched_thresh') sign_peaks = params.get('detection', 'peaks') do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = params.getint('whitening', 'chunk_size') plot_path = os.path.join(params.get('data', 'data_file_noext'), 'plots') nodes, edges = get_nodes_and_edges(params) safety_time = params.getint('whitening', 'safety_time') nb_temp_white = min(max(20, comm.size), N_e) max_silence_1 = int(20 * params.rate // comm.size) max_silence_2 = 5000 inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.argsort(nodes) ################################################################# if comm.rank == 0: print_and_log( ["Analyzing data to get whitening matrices and thresholds..."], 'default', logger) if use_gpu: import cudamat as cmt ## Need to properly handle multi GPU per MPI nodes? if nb_gpu > nb_cpu: gpu_id = int(comm.rank // nb_cpu) else: gpu_id = 0 cmt.cuda_set_device(gpu_id) cmt.init() cmt.cuda_sync_threads() nb_chunks, last_chunk_len = data_file.analyze(chunk_size) if nb_chunks < comm.size: res = io.data_stats(params, show=False) chunk_size = int(res * params.rate // comm.size) if comm.rank == 0: print_and_log( ["Too much cores, automatically resizing the data chunks"], 'debug', logger) nb_chunks, last_chunk_len = data_file.analyze(chunk_size) # I guess this is more relevant, to take signals from all over the recordings all_chunks = numpy.random.permutation( numpy.arange(nb_chunks, dtype=numpy.int32)) all_electrodes = numpy.random.permutation(N_e) for gidx in [all_chunks[comm.rank]]: #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." local_chunk, t_offset = data_file.get_data(gidx, chunk_size, nodes=nodes) local_shape = len(local_chunk) #print "Node", comm.rank, "computes the median absolute deviations in a random chunk" thresholds = numpy.zeros(N_e, dtype=numpy.float32) for i in xrange(N_e): u = numpy.median(local_chunk[:, i], 0) thresholds[i] = numpy.median(numpy.abs(local_chunk[:, i] - u), 0) gdata = gather_array(thresholds, comm) if comm.rank == 0: gdata = gdata.reshape((comm.size, N_e)) thresholds = numpy.mean(gdata, 0) bfile = h5py.File(file_out_suff + '.basis.hdf5', 'w', libver='latest') io.write_datasets(bfile, ['thresholds'], {'thresholds': thresholds}) bfile.close() comm.Barrier() thresholds = io.load_data(params, 'thresholds') #print "Extracting the peaks..." local_peaktimes = numpy.zeros(0, dtype=numpy.int32) for i in xrange(N_e): peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False, mpd=dist_peaks) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) local_peaktimes = numpy.unique(local_peaktimes) #print "Removing the useless borders..." local_borders = (template_shift, local_shape - template_shift) idx = (local_peaktimes >= local_borders[0]) & (local_peaktimes < local_borders[1]) local_peaktimes = numpy.compress(idx, local_peaktimes) if len(local_peaktimes) > 0: diff_times = local_peaktimes[-1] - local_peaktimes[0] all_times = numpy.zeros((N_e, diff_times + 1), dtype=numpy.bool) min_times = numpy.maximum( local_peaktimes - local_peaktimes[0] - safety_time, 0) max_times = numpy.minimum( local_peaktimes - local_peaktimes[0] + safety_time + 1, diff_times) argmax_peak = numpy.random.permutation( numpy.arange(len(local_peaktimes))) all_idx = numpy.take(local_peaktimes, argmax_peak) #print "Selection of the peaks with spatio-temporal masks..." for idx, peak in zip(argmax_peak, all_idx): elec = numpy.argmax(numpy.abs(local_chunk[peak])) indices = numpy.take(inv_nodes, edges[nodes[elec]]) all_times[indices, min_times[idx]:max_times[idx]] = True else: all_times = numpy.zeros((N_e, len(local_chunk)), dtype=numpy.bool) all_times_Ne = numpy.any(all_times, 0) subset = numpy.where(all_times_Ne == False)[0] all_silences = [] if do_spatial_whitening: local_silences = numpy.take(local_chunk, subset, axis=0)[:max_silence_1] all_silences = gather_array(local_silences, comm, 0, 1) local_res = [] if do_temporal_whitening: for elec in all_electrodes[numpy.arange(comm.rank, nb_temp_white, comm.size)]: res = numpy.zeros((0, N_t), dtype=numpy.float32) scount = 0 indices = numpy.take(inv_nodes, edges[nodes[elec]]) all_times_elec = numpy.any(numpy.take(all_times, indices, axis=0), 0) esubset = numpy.where(all_times_elec == False)[0] bound = len(esubset) - N_t while (scount < bound) and (len(res) < max_silence_2): myslice = esubset[scount:scount + N_t] if numpy.all((myslice - esubset[scount]) == numpy.arange(N_t)): scount += N_t res = numpy.vstack((res, local_chunk[myslice, elec])) else: scount += 1 if len(res) > 5: local_res += [numpy.cov(res.T)] nb_elecs = numpy.array([len(local_res)], dtype=numpy.float32) local_res = numpy.array(local_res, dtype=numpy.float32) if len(local_res) == 0: local_res = numpy.zeros(0, dtype=numpy.float32) else: local_res = numpy.sum(local_res, 0) all_res = gather_array(local_res.ravel(), comm, 0, 1) all_elecs = gather_array(nb_elecs, comm, 0, 1) if comm.rank == 0: to_write = {} if do_temporal_whitening: try: nb_silences = numpy.sum(all_elecs > 0) all_res = all_res.reshape((nb_silences, N_t**2)) except Exception: print_and_log([ "No silent periods detected: something wrong with the parameters?" ], 'error', logger) all_res = numpy.sum(all_res, 0) all_res = all_res.reshape((N_t, N_t)) / numpy.sum(all_elecs) temporal_whitening = get_whitening_matrix( all_res.astype(numpy.double), fudge=1e-3)[template_shift].astype(numpy.float32) temporal_whitening /= temporal_whitening.sum() to_write['temporal'] = temporal_whitening have_nans = numpy.sum(numpy.isnan(temporal_whitening)) if have_nans > 0: temporal_whitening = numpy.zeros(N_t, dtype=numpy.float32) temporal_whitening[N_t // 2] = 1 to_write['temporal'] = temporal_whitening print_and_log( ["Disabling temporal whitening because of NaNs found"], 'info', logger) if do_spatial_whitening: if len(all_silences) / params.rate == 0: print_and_log([ "No silent periods detected: something wrong with the parameters?" ], 'error', logger) spatial_whitening = get_whitening_matrix( all_silences.astype(numpy.double)).astype(numpy.float32) to_write['spatial'] = spatial_whitening print_and_log([ "Found %gs without spikes for whitening matrices..." % (len(all_silences) / params.rate) ], 'default', logger) have_nans = numpy.sum(numpy.isnan(spatial_whitening)) if have_nans > 0: spatial_whitening = numpy.eye(spatial_whitening.shape[0], dtype=numpy.float32) to_write['spatial'] = spatial_whitening print_and_log( ["Disabling spatial whitening because of NaNs found"], 'info', logger) bfile = h5py.File(file_out_suff + '.basis.hdf5', 'r+', libver='latest') io.write_datasets(bfile, to_write.keys(), to_write) bfile.close() del all_silences comm.Barrier() if do_spatial_whitening or do_temporal_whitening: if comm.rank == 0: print_and_log( ["Because of whitening, need to recompute the thresholds..."], 'default', logger) if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if use_gpu: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') for gidx in [all_chunks[comm.rank]]: local_chunk, t_offset = data_file.get_data(gidx, chunk_size, nodes=nodes) local_shape = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d( local_chunk, temporal_whitening, axis=0, mode='constant') thresholds = numpy.zeros(N_e, dtype=numpy.float32) for i in xrange(N_e): u = numpy.median(local_chunk[:, i], 0) thresholds[i] = numpy.median(numpy.abs(local_chunk[:, i] - u), 0) gdata = gather_array(thresholds, comm) if comm.rank == 0: gdata = gdata.reshape((comm.size, N_e)) thresholds = numpy.mean(gdata, 0) bfile = h5py.File(file_out_suff + '.basis.hdf5', 'r+', libver='latest') bfile.pop('thresholds') io.write_datasets(bfile, ['thresholds'], {'thresholds': thresholds}) bfile.close() comm.Barrier() #if comm.rank == 0: #if not os.path.exists(plot_path): # os.makedirs(plot_path) #N_elec = min(int(numpy.sqrt(data_file.N_e)), 5) #plot.view_fit(filename, t_start=0, t_stop=1, fit_on=False, square=True, # n_elec=N_elec, save=[plot_path, 'electrodes']) # Part 2: Basis numpy.random.seed(422) ################################################################# file_out = params.get('data', 'file_out') alignment = params.getboolean('detection', 'alignment') spike_thresh = params.getfloat('detection', 'spike_thresh') nodes, edges = get_nodes_and_edges(params) do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = params.getint('data', 'chunk_size') safety_time = params.getint('whitening', 'safety_time') max_elts_elec = params.getint('whitening', 'max_elts') nb_elts = int( params.getfloat('whitening', 'nb_elts') * N_e * max_elts_elec) output_dim = params.getfloat('whitening', 'output_dim') inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.argsort(nodes) if sign_peaks == 'both': max_elts_elec *= 2 ################################################################# if comm.rank == 0: print_and_log(["Searching spikes to construct the PCA basis..."], 'default', logger) nb_chunks, last_chunk_len = data_file.analyze(chunk_size) if nb_chunks < comm.size: res = io.data_stats(params, show=False) chunk_size = int(res * params.rate // comm.size) if comm.rank == 0: print_and_log( ["Too much cores, automatically resizing the data chunks"], 'debug', logger) nb_chunks, last_chunk_len = data_file.analyze(chunk_size) groups = {} for i in xrange(N_e): groups[i] = 0 # I guess this is more relevant, to take signals from all over the recordings all_chunks = numpy.random.permutation( numpy.arange(nb_chunks, dtype=numpy.int32)) max_elts_elec //= comm.size nb_elts //= comm.size elt_count_pos = 0 elt_count_neg = 0 if sign_peaks in ['positive', 'both']: elts_pos = numpy.zeros((N_t, nb_elts), dtype=numpy.float32) if sign_peaks in ['negative', 'both']: elts_neg = numpy.zeros((N_t, nb_elts), dtype=numpy.float32) chunks_to_load = all_chunks[comm.rank::comm.size] thresholds = io.load_data(params, 'thresholds') if comm.rank == 0: pbar = get_progressbar(nb_elts) if alignment: cdata = numpy.linspace(-template_shift, template_shift, 5 * N_t) xdata = numpy.arange(-2 * template_shift, 2 * template_shift + 1) for gcount, gidx in enumerate(chunks_to_load): if ((elt_count_pos + elt_count_neg) < nb_elts): #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." local_chunk, t_offset = data_file.get_data(gidx, chunk_size, nodes=nodes) local_shape = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d( local_chunk, temporal_whitening, axis=0, mode='constant') #print "Extracting the peaks..." all_peaktimes = numpy.zeros(0, dtype=numpy.int32) all_extremas = numpy.zeros(0, dtype=numpy.int32) for i in xrange(N_e): if sign_peaks == 'negative': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=True, mpd=dist_peaks) elif sign_peaks == 'positive': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=False, mpd=dist_peaks) elif sign_peaks == 'both': peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False, mpd=dist_peaks) all_peaktimes = numpy.concatenate((all_peaktimes, peaktimes)) all_extremas = numpy.concatenate( (all_extremas, i * numpy.ones(len(peaktimes), dtype=numpy.int32))) #print "Removing the useless borders..." if alignment: local_borders = (2 * template_shift, local_shape - 2 * template_shift) else: local_borders = (template_shift, local_shape - template_shift) idx = (all_peaktimes >= local_borders[0]) & (all_peaktimes < local_borders[1]) all_peaktimes = numpy.compress(idx, all_peaktimes) all_extremas = numpy.compress(idx, all_extremas) local_peaktimes = numpy.unique(all_peaktimes) if len(local_peaktimes) > 0: diff_times = local_peaktimes[-1] - local_peaktimes[0] all_times = numpy.zeros((N_e, diff_times + 1), dtype=numpy.bool) min_times = numpy.maximum( local_peaktimes - local_peaktimes[0] - safety_time, 0) max_times = numpy.minimum( local_peaktimes - local_peaktimes[0] + safety_time + 1, diff_times) n_times = len(local_peaktimes) argmax_peak = numpy.random.permutation(numpy.arange(n_times)) all_idx = numpy.take(local_peaktimes, argmax_peak) #print "Selection of the peaks with spatio-temporal masks..." for midx, peak in zip(argmax_peak, all_idx): if (elt_count_neg + elt_count_pos) == nb_elts: break if sign_peaks == 'negative': elec = numpy.argmin(local_chunk[peak]) negative_peak = True elif sign_peaks == 'positive': elec = numpy.argmax(local_chunk[peak]) negative_peak = False elif sign_peaks == 'both': if numpy.abs(numpy.max(local_chunk[peak])) > numpy.abs( numpy.min(local_chunk[peak])): elec = numpy.argmax(local_chunk[peak]) negative_peak = False else: elec = numpy.argmin(local_chunk[peak]) negative_peak = True indices = numpy.take(inv_nodes, edges[nodes[elec]]) myslice = all_times[indices, min_times[midx]:max_times[midx]] is_local_extrema = elec in all_extremas[all_peaktimes == peak] if is_local_extrema and not myslice.any(): upper_bounds = max_elts_elec if groups[elec] < upper_bounds: if negative_peak: elts_neg[:, elt_count_neg] = local_chunk[ peak - template_shift:peak + template_shift + 1, elec] else: elts_pos[:, elt_count_pos] = local_chunk[ peak - template_shift:peak + template_shift + 1, elec] if alignment: ydata = local_chunk[peak - 2 * template_shift:peak + 2 * template_shift + 1, elec] f = scipy.interpolate.UnivariateSpline(xdata, ydata, s=0) if negative_peak: rmin = (numpy.argmin(f(cdata)) - len(cdata) / 2.) / 5. else: rmin = (numpy.argmax(f(cdata)) - len(cdata) / 2.) / 5. ddata = numpy.linspace(rmin - template_shift, rmin + template_shift, N_t) if negative_peak: elts_neg[:, elt_count_neg] = f(ddata).astype( numpy.float32) else: elts_pos[:, elt_count_pos] = f(ddata).astype( numpy.float32) if negative_peak: elt_count_neg += 1 else: elt_count_pos += 1 groups[elec] += 1 all_times[indices, min_times[midx]:max_times[midx]] = True if comm.rank == 0: pbar.update(elt_count_pos + elt_count_neg) if comm.rank == 0: if (elt_count_pos + elt_count_neg < (gcount + 1) * max_elts_elec // len(chunks_to_load)): pbar.update( (gcount + 1) * max_elts_elec // len(chunks_to_load)) if comm.rank == 0: pbar.finish() print_and_log([ "Node %d has collected %d waveforms" % (comm.rank, elt_count_pos + elt_count_neg) ], 'debug', logger) if sign_peaks in ['negative', 'both']: gdata_neg = gather_array(elts_neg[:, :elt_count_neg].T, comm, 0, 1) if sign_peaks in ['positive', 'both']: gdata_pos = gather_array(elts_pos[:, :elt_count_pos].T, comm, 0, 1) if comm.rank == 0: #DO PCA on elts and store the basis obtained. nb_waveforms = 0 if sign_peaks in ['negative', 'both']: nb_waveforms += gdata_neg.shape[0] if sign_peaks in ['positive', 'both']: nb_waveforms += gdata_pos.shape[0] print_and_log([ "Found %d waveforms over %d requested" % (nb_waveforms, int(nb_elts * comm.size)) ], 'default', logger) pca = PCA(output_dim, copy=False) res = {} if sign_peaks in ['negative', 'both']: if len(gdata_neg) > 0: res_pca = pca.fit_transform(gdata_neg.astype( numpy.double)).astype(numpy.float32) res['proj'] = pca.components_.T.astype(numpy.float32) else: res['proj'] = numpy.identity(N_t, dtype=numpy.float32) res['rec'] = res['proj'].T res['waveform'] = numpy.median(gdata_neg, 0) idx = numpy.random.permutation(numpy.arange( gdata_neg.shape[0]))[:1000] res['waveforms'] = gdata_neg[idx, :] if sign_peaks in ['positive', 'both']: if len(gdata_pos) > 0: res_pca = pca.fit_transform(gdata_pos.astype( numpy.double)).astype(numpy.float32) res['proj_pos'] = pca.components_.T.astype(numpy.float32) else: res['proj_pos'] = numpy.identity(N_t, dtype=numpy.float32) res['rec_pos'] = res['proj_pos'].T res['waveform_pos'] = numpy.median(gdata_pos, 0) idx = numpy.random.permutation(numpy.arange( gdata_pos.shape[0]))[:1000] res['waveforms_pos'] = gdata_pos[idx, :] bfile = h5py.File(file_out_suff + '.basis.hdf5', 'r+', libver='latest') io.write_datasets(bfile, res.keys(), res) if sign_peaks == 'positive': print_and_log([ "A basis with %s dimensions has been built" % res['proj_pos'].shape[1] ], 'info', logger) elif sign_peaks == 'negative': print_and_log([ "A basis with %s dimensions has been built" % res['proj'].shape[1] ], 'info', logger) elif sign_peaks == 'both': print_and_log([ "Two basis with %s dimensions has been built" % res['proj'].shape[1] ], 'info', logger) bfile.close() comm.Barrier() if matched_filter: if comm.rank == 0: print_and_log([ "Because of matched filters, need to recompute the thresholds..." ], 'default', logger) if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if use_gpu: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') if sign_peaks in ['negative', 'both']: waveform_neg = io.load_data(params, 'waveform') waveform_neg /= (numpy.abs(numpy.sum(waveform_neg)) * len(waveform_neg)) if sign_peaks in ['positive', 'both']: waveform_pos = io.load_data(params, 'waveform-pos') waveform_pos /= (numpy.abs(numpy.sum(waveform_pos)) * len(waveform_pos)) for gidx in [all_chunks[comm.rank]]: local_chunk, t_offset = data_file.get_data(gidx, chunk_size, nodes=nodes) local_shape = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d( local_chunk, temporal_whitening, axis=0, mode='constant') if sign_peaks in ['negative', 'both']: tmp_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_neg, axis=0, mode='constant') thresholds = numpy.zeros(N_e, dtype=numpy.float32) for i in xrange(N_e): u = numpy.median(tmp_chunk[:, i], 0) thresholds[i] = numpy.median( numpy.abs(tmp_chunk[:, i] - u), 0) gdata = gather_array(thresholds, comm) if comm.rank == 0: gdata = gdata.reshape((comm.size, N_e)) thresholds = numpy.mean(gdata, 0) bfile = h5py.File(file_out + '.basis.hdf5', 'r+', libver='latest') io.write_datasets(bfile, ['matched_thresholds'], {'matched_thresholds': thresholds}) bfile.close() comm.Barrier() if sign_peaks in ['positive', 'both']: tmp_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_pos, axis=0, mode='constant') thresholds = numpy.zeros(N_e, dtype=numpy.float32) for i in xrange(N_e): u = numpy.median(tmp_chunk[:, i], 0) thresholds[i] = numpy.median( numpy.abs(tmp_chunk[:, i] - u), 0) gdata = gather_array(thresholds, comm) if comm.rank == 0: gdata = gdata.reshape((comm.size, N_e)) thresholds = numpy.mean(gdata, 0) bfile = h5py.File(file_out + '.basis.hdf5', 'r+', libver='latest') io.write_datasets(bfile, ['matched_thresholds_pos'], {'matched_thresholds_pos': thresholds}) bfile.close() comm.Barrier() data_file.close()
def main(params, nb_cpu, nb_gpu, use_gpu): ################################################################# #params = detect_memory(params) logger = init_logging(params.logfile) SHARED_MEMORY = get_shared_memory_flag(params) logger = logging.getLogger('circus.fitting') data_file = params.data_file data_file.open() N_e = params.getint('data', 'N_e') N_total = params.nb_channels N_t = params.getint('detection', 'N_t') template_shift = params.getint('detection', 'template_shift') file_out = params.get('data', 'file_out') file_out_suff = params.get('data', 'file_out_suff') sign_peaks = params.get('detection', 'peaks') matched_filter = params.getboolean('detection', 'matched-filter') spike_thresh = params.getfloat('detection', 'spike_thresh') do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = params.getint('fitting', 'chunk_size') gpu_only = params.getboolean('fitting', 'gpu_only') nodes, edges = get_nodes_and_edges(params) tmp_limits = params.get('fitting', 'amp_limits').replace('(', '').replace(')', '').split(',') tmp_limits = map(float, tmp_limits) amp_auto = params.getboolean('fitting', 'amp_auto') nb_chances = params.getint('fitting', 'nb_chances') max_chunk = params.getfloat('fitting', 'max_chunk') noise_thr = params.getfloat('clustering', 'noise_thr') collect_all = params.getboolean('fitting', 'collect_all') ignore_dead_times = params.getboolean('triggers', 'ignore_times') inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.argsort(nodes) ################################################################# if use_gpu: import cudamat as cmt ## Need to properly handle multi GPU per MPI nodes? if nb_gpu > nb_cpu: gpu_id = int(comm.rank // nb_cpu) else: gpu_id = 0 cmt.cuda_set_device(gpu_id) cmt.init() cmt.cuda_sync_threads() if SHARED_MEMORY: templates = io.load_data_memshared(params, 'templates', normalize=True, transpose=True) N_tm, x = templates.shape else: templates = io.load_data(params, 'templates') x, N_tm = templates.shape temp_2_shift = 2 * template_shift full_gpu = use_gpu and gpu_only n_tm = N_tm // 2 n_scalar = N_e * N_t temp_window = numpy.arange(-template_shift, template_shift + 1) size_window = N_e * (2 * template_shift + 1) if not amp_auto: amp_limits = numpy.zeros((n_tm, 2)) amp_limits[:, 0] = tmp_limits[0] amp_limits[:, 1] = tmp_limits[1] else: amp_limits = io.load_data(params, 'limits') norm_templates = io.load_data(params, 'norm-templates') if not SHARED_MEMORY: for idx in xrange(templates.shape[1]): myslice = numpy.arange(templates.indptr[idx], templates.indptr[idx + 1]) templates.data[myslice] /= norm_templates[idx] templates = templates.T if matched_filter: if sign_peaks in ['negative', 'both']: waveform_neg = io.load_data(params, 'waveform') waveform_neg /= (numpy.abs(numpy.sum(waveform_neg)) * len(waveform_neg)) matched_tresholds_neg = io.load_data(params, 'matched-thresholds') if sign_peaks in ['positive', 'both']: waveform_pos = io.load_data(params, 'waveform-pos') waveform_pos /= (numpy.abs(numpy.sum(waveform_pos)) * len(waveform_pos)) matched_tresholds_pos = io.load_data(params, 'matched-thresholds-pos') if ignore_dead_times: all_dead_times = get_dead_times(params) thresholds = io.load_data(params, 'thresholds') if collect_all: neighbors = {} for i in xrange(n_tm): tmp = templates[i, :].toarray().reshape(N_e, N_t) * norm_templates[i] neighbors[i] = numpy.where(numpy.sum(tmp, 1) != 0)[0] if use_gpu: templates = cmt.SparseCUDAMatrix(templates, copy_on_host=False) info_string = '' if comm.rank == 0: if use_gpu: info_string = "using %d GPUs" % (comm.size) else: info_string = "using %d CPUs" % (comm.size) comm.Barrier() c_overlap = io.get_overlaps(params, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] ## If the number of overlaps is different from templates, we need to recompute them if N_over != N_tm: if comm.rank == 0: print_and_log( ['Templates have been modified, recomputing the overlaps...'], 'default', logger) c_overlap = io.get_overlaps(params, erase=True, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] if SHARED_MEMORY: c_overs = io.load_data_memshared(params, 'overlaps') else: c_overs = io.load_data(params, 'overlaps') comm.Barrier() if n_tm == 0: if comm.rank == 0: print_and_log(["No templates present. Redo clustering?"], 'default', logger) sys.exit(0) if comm.rank == 0: print_and_log([ "Here comes the SpyKING CIRCUS %s and %d templates..." % (info_string, n_tm) ], 'default', logger) purge(file_out_suff, '.data') if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') if full_gpu: try: # If memory on the GPU is large enough, we load the overlaps onto it for i in xrange(N_over): c_overs[i] = cmt.SparseCUDAMatrix(c_overs[i], copy_on_host=False) except Exception: if comm.rank == 0: print_and_log([ "Not enough memory on GPUs: GPUs are used for projection only" ], 'info', logger) for i in xrange(N_over): if c_overs.has_key(i): del c_overs[i] full_gpu = False nb_chunks, last_chunk_len = data_file.analyze(chunk_size) processed_chunks = int(min(nb_chunks, max_chunk)) comm.Barrier() spiketimes_file = open(file_out_suff + '.spiketimes-%d.data' % comm.rank, 'wb') comm.Barrier() amplitudes_file = open(file_out_suff + '.amplitudes-%d.data' % comm.rank, 'wb') comm.Barrier() templates_file = open(file_out_suff + '.templates-%d.data' % comm.rank, 'wb') comm.Barrier() if collect_all: garbage_times_file = open( file_out_suff + '.gspiketimes-%d.data' % comm.rank, 'wb') comm.Barrier() garbage_temp_file = open( file_out_suff + '.gtemplates-%d.data' % comm.rank, 'wb') comm.Barrier() if use_gpu and do_spatial_whitening: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) last_chunk_size = 0 to_explore = xrange(comm.rank, processed_chunks, comm.size) if comm.rank == 0: to_explore = get_tqdm_progressbar(to_explore) for gcount, gidx in enumerate(to_explore): #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." ## We need to deal with the borders by taking chunks of size [0, chunck_size+template_shift] is_first = data_file.is_first_chunk(gidx, nb_chunks) is_last = data_file.is_last_chunk(gidx, nb_chunks) if is_last: padding = (-temp_2_shift, 0) elif is_first: padding = (0, temp_2_shift) else: padding = (-temp_2_shift, temp_2_shift) result = {'spiketimes': [], 'amplitudes': [], 'templates': []} local_chunk, t_offset = data_file.get_data(gidx, chunk_size, padding, nodes=nodes) len_chunk = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d(local_chunk, temporal_whitening, axis=0, mode='constant') #print "Extracting the peaks..." if collect_all: all_found_spikes = {} for i in xrange(N_e): all_found_spikes[i] = [] local_peaktimes = numpy.zeros(0, dtype=numpy.uint32) if matched_filter: if sign_peaks in ['positive', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d( local_chunk, waveform_pos, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_pos[i]) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() if sign_peaks in ['negative', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d( local_chunk, waveform_neg, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_neg[i]) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() else: for i in xrange(N_e): if sign_peaks == 'negative': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=True) elif sign_peaks == 'positive': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=False) elif sign_peaks == 'both': peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() local_peaktimes = numpy.unique(local_peaktimes) g_offset = t_offset + padding[0] if ignore_dead_times: dead_indices = numpy.searchsorted( all_dead_times, [t_offset, t_offset + chunk_size]) if dead_indices[0] != dead_indices[1]: local_peaktimes = numpy.array(list( set(local_peaktimes + g_offset).difference( all_dead_times[dead_indices[0]:dead_indices[1]])), dtype=numpy.uint32) - g_offset local_peaktimes = numpy.sort(local_peaktimes) #print "Removing the useless borders..." local_borders = (template_shift, len_chunk - template_shift) idx = (local_peaktimes >= local_borders[0]) & (local_peaktimes < local_borders[1]) local_peaktimes = numpy.compress(idx, local_peaktimes) if collect_all: for i in xrange(N_e): all_found_spikes[i] = numpy.array(all_found_spikes[i], dtype=numpy.uint32) if ignore_dead_times: if dead_indices[0] != dead_indices[1]: all_found_spikes[i] = numpy.array( list( set(all_found_spikes[i] + g_offset).difference( all_dead_times[ dead_indices[0]:dead_indices[1]])), dtype=numpy.uint32) - g_offset all_found_spikes[i] = numpy.sort(all_found_spikes[i]) idx = (all_found_spikes[i] >= local_borders[0]) & ( all_found_spikes[i] < local_borders[1]) all_found_spikes[i] = numpy.compress(idx, all_found_spikes[i]) n_t = len(local_peaktimes) if full_gpu: # all_indices = cmt.CUDAMatrix(all_indices) tmp_gpu = cmt.CUDAMatrix(local_peaktimes.reshape((1, n_t)), copy_on_host=False) if n_t > 0: #print "Computing the b (should full_gpu by putting all chunks on GPU if possible?)..." if collect_all: c_local_chunk = local_chunk.copy() local_chunk = local_chunk.T.ravel() sub_mat = numpy.zeros((size_window, n_t), dtype=numpy.float32) if len_chunk != last_chunk_size: slice_indices = numpy.zeros(0, dtype=numpy.int32) for idx in xrange(N_e): slice_indices = numpy.concatenate( (slice_indices, len_chunk * idx + temp_window)) last_chunk_size = len_chunk for count, idx in enumerate(local_peaktimes): sub_mat[:, count] = numpy.take(local_chunk, slice_indices + idx) del local_chunk if use_gpu: sub_mat = cmt.CUDAMatrix(sub_mat, copy_on_host=False) b = cmt.sparse_dot(templates, sub_mat) else: b = templates.dot(sub_mat) del sub_mat local_restriction = (t_offset, t_offset + chunk_size) all_spikes = local_peaktimes + g_offset # Because for GPU, slicing by columns is more efficient, we need to transpose b #b = b.transpose() if use_gpu and not full_gpu: b = b.asarray() failure = numpy.zeros(n_t, dtype=numpy.int32) if full_gpu: mask = numpy.zeros((2 * n_tm, n_t), dtype=numpy.float32) mask[:n_tm, :] = 1 data = cmt.empty(mask.shape) patch_gpu = b.shape[1] == 1 else: mask = numpy.ones((n_tm, n_t), dtype=numpy.float32) sub_b = b[:n_tm, :] if collect_all: c_all_times = numpy.zeros((len_chunk, N_e), dtype=numpy.bool) c_min_times = numpy.maximum( numpy.arange(len_chunk) - template_shift, 0) c_max_times = numpy.minimum( numpy.arange(len_chunk) + template_shift + 1, len_chunk) for i in xrange(N_e): c_all_times[all_found_spikes[i], i] = True while (numpy.mean(failure) < nb_chances): if full_gpu: gpu_mask = cmt.CUDAMatrix(mask, copy_on_host=False) b.mult(gpu_mask, data) tmp_mat = data.max(0) argmax_bi = numpy.argsort(tmp_mat.asarray()[0, :])[::-1] del tmp_mat else: data = sub_b * mask argmax_bi = numpy.argsort(numpy.max(data, 0))[::-1] for peak_index in argmax_bi: if full_gpu: b_array = b.asarray() sub_b = b_array[:n_tm, :] peak_scalar_products = np.take(sub_b, peak_index, axis=1) best_template_index = np.argmax(peak_scalar_products, axis=0) best_template2_index = best_template_index + n_tm if full_gpu: best_amp = sub_b[best_template_index, peak_index] / n_scalar best_amp2 = b_array[best_template_index, peak_index] / n_scalar else: best_amp = sub_b[best_template_index, peak_index] / n_scalar best_amp2 = b[best_template2_index, peak_index] / n_scalar best_amp_n = best_amp / norm_templates[best_template_index] best_amp2_n = best_amp2 / norm_templates[ best_template2_index] # Verify amplitude constraint. a_min = amp_limits[best_template_index, 0] a_max = amp_limits[best_template_index, 1] if (a_min <= best_amp_n) & (best_amp_n <= a_max): # Keep the matching. peak_time_step = local_peaktimes[peak_index] data = (local_peaktimes - peak_time_step).astype( np.int32) is_neighbor = np.where(np.abs(data) <= temp_2_shift)[0] idx_neighbor = data[is_neighbor] + temp_2_shift nb_neighbors = len(is_neighbor) indices = np.zeros((S_over, nb_neighbors), dtype=np.int32) indices[idx_neighbor, np.arange(nb_neighbors)] = 1 if full_gpu: indices = cmt.CUDAMatrix(indices, copy_on_host=False) if patch_gpu: b_lines = b.get_col_slice(0, b.shape[0]) else: b_lines = b.get_col_slice( is_neighbor[0], is_neighbor[-1] + 1) tmp1 = cmt.sparse_dot(c_overs[best_template_index], indices, mult=-best_amp[keep]) tmp2 = cmt.sparse_dot( c_overs[best_template2_index], indices, mult=-best_amp2[keep]) b_lines.add(tmp1.add(tmp2)) del tmp1, tmp2 else: tmp1 = c_overs[best_template_index].multiply( -best_amp) tmp2 = c_overs[best_template2_index].multiply( -best_amp2) b[:, is_neighbor] += (tmp1 + tmp2).dot(indices) # Add matching to the result. t_spike = all_spikes[peak_index] if (t_spike >= local_restriction[0]) and ( t_spike < local_restriction[1]): #print "Accept spikes", t_spike, local_restriction, type(t_spike), t_spike > local_restriction[0], t_spike < local_restriction[1] result['spiketimes'] += [t_spike] result['amplitudes'] += [(best_amp_n, best_amp2_n)] result['templates'] += [best_template_index] # Mark current matching as tried. mask[best_template_index, peak_index] = 0 else: # Reject the matching. # Update failure counter of the peak. failure[peak_index] += 1 # If the maximal number of failures is reached then mark peak as solved (i.e. not fitted). if failure[peak_index] == nb_chances: mask[:, peak_index] = 0 else: mask[best_template_index, peak_index] = 0 spikes_to_write = numpy.array(result['spiketimes'], dtype=numpy.uint32) amplitudes_to_write = numpy.array(result['amplitudes'], dtype=numpy.float32) templates_to_write = numpy.array(result['templates'], dtype=numpy.uint32) spiketimes_file.write(spikes_to_write.tostring()) amplitudes_file.write(amplitudes_to_write.tostring()) templates_file.write(templates_to_write.tostring()) if collect_all: for temp, spike in zip(templates_to_write, spikes_to_write - g_offset): c_all_times[c_min_times[spike]:c_max_times[spike], neighbors[temp]] = False gspikes = numpy.where(numpy.sum(c_all_times, 1) > 0)[0] c_all_times = numpy.take(c_all_times, gspikes, axis=0) c_local_chunk = numpy.take(c_local_chunk, gspikes, axis=0) * c_all_times if sign_peaks == 'negative': bestlecs = numpy.argmin(c_local_chunk, 1) if matched_filter: threshs = -matched_tresholds_neg[bestlecs] else: threshs = -thresholds[bestlecs] idx = numpy.where(numpy.min(c_local_chunk, 1) < threshs)[0] elif sign_peaks == 'positive': bestlecs = numpy.argmax(c_local_chunk, 1) if matched_filter: threshs = matched_tresholds_pos[bestlecs] else: threshs = thresholds[bestlecs] idx = numpy.where(numpy.max(c_local_chunk, 1) > threshs)[0] elif sign_peaks == 'both': c_local_chunk = numpy.abs(c_local_chunk) bestlecs = numpy.argmax(c_local_chunk, 1) if matched_filter: threshs = numpy.minimum( matched_tresholds_neg[bestlecs], matched_tresholds_pos[bestlecs]) else: threshs = thresholds[bestlecs] idx = numpy.where(numpy.max(c_local_chunk, 1) > threshs)[0] gspikes = numpy.take(gspikes, idx) bestlecs = numpy.take(bestlecs, idx) gspikes_to_write = numpy.array(gspikes + g_offset, dtype=numpy.uint32) gtemplates_to_write = numpy.array(bestlecs, dtype=numpy.uint32) garbage_times_file.write(gspikes_to_write.tostring()) garbage_temp_file.write(gtemplates_to_write.tostring()) if full_gpu: del gpu_mask, b, data sys.stderr.flush() spiketimes_file.flush() os.fsync(spiketimes_file.fileno()) spiketimes_file.close() amplitudes_file.flush() os.fsync(amplitudes_file.fileno()) amplitudes_file.close() templates_file.flush() os.fsync(templates_file.fileno()) templates_file.close() if collect_all: garbage_temp_file.flush() os.fsync(garbage_temp_file.fileno()) garbage_temp_file.close() garbage_times_file.flush() os.fsync(garbage_times_file.fileno()) garbage_times_file.close() comm.Barrier() if comm.rank == 0: io.collect_data(comm.size, params, erase=True) data_file.close()
def extract_juxta_spikes_(params): '''Detect spikes from the extracellular traces''' file_out_suff = params.get('data', 'file_out_suff') sampling_rate = params.getint('data', 'sampling_rate') dist_peaks = params.getint('detection', 'dist_peaks') template_shift = params.getint('detection', 'template_shift') juxta_dtype = params.get('validating', 'juxta_dtype') juxta_thresh = params.getfloat('validating', 'juxta_thresh') juxta_valley = params.getboolean('validating', 'juxta_valley') juxta_spikes = params.get('validating', 'juxta_spikes') juxta_filename = "{}.juxta.dat".format(file_out_suff) beer_path = "{}.beer.hdf5".format(file_out_suff) if juxta_spikes == '': # Read juxtacellular trace. juxta_data = numpy.fromfile(juxta_filename, dtype=juxta_dtype) #juxta_data = juxta_data.astype(numpy.float32) # juxta_data = juxta_data - dtype_offset juxta_data = numpy.ascontiguousarray(juxta_data) # Filter juxtacellular trace. juxta_data = highpass(juxta_data, sampling_rate=sampling_rate) juxta_data -= numpy.median(juxta_data) # Compute median and median absolute deviation. juxta_median = numpy.median(juxta_data) juxta_ad = numpy.abs(juxta_data - juxta_median) juxta_mad = numpy.median(juxta_ad, axis=0) # Save medians and median absolute deviations to BEER file. beer_file = h5py.File(beer_path, 'a', libver='latest') if "juxta_median" in beer_file.keys(): beer_file.pop("juxta_median") beer_file.create_dataset("juxta_median", data=juxta_median) if "juxta_mad" in beer_file.keys(): beer_file.pop("juxta_mad") beer_file.create_dataset("juxta_mad", data=juxta_mad) beer_file.close() if comm.rank == 0: print_and_log(["Extract juxtacellular spikes"], level='debug', logger=logger) # Detect juxta spike times. threshold = juxta_thresh * juxta_mad juxta_spike_times = algo.detect_peaks(juxta_data, threshold, valley=juxta_valley, mpd=dist_peaks) # Remove juxta spike times in the borders. juxta_spike_times = juxta_spike_times[ template_shift <= juxta_spike_times] juxta_spike_times = juxta_spike_times[ juxta_spike_times < juxta_data.size - template_shift] else: juxta_spike_times = numpy.load(juxta_spikes) # Save juxta spike times to BEER file. beer_file = h5py.File(beer_path, 'a', libver='latest') group_name = "juxta_spiketimes" if group_name in beer_file.keys(): beer_file.pop(group_name) beer_file.create_group(group_name) key = "{}/elec_0".format(group_name) beer_file.create_dataset(key, data=juxta_spike_times) beer_file.close() # juxta_spike_values = numpy.zeros_like(juxta_spike_times, dtype='float') # for i, t in enumerate(juxta_spike_times): # if juxta_valley: # juxta_spike_values[i] = - juxta_data[t] # else: # juxta_spike_values[i] = + juxta_data[t] if juxta_spikes == '': # Find juxta spike values of juxta spike times. juxta_spike_values = juxta_data[juxta_spike_times] if juxta_valley: juxta_spike_values *= -1 # Save juxta spike values to BEER file. beer_file = h5py.File(beer_path, 'a', libver='latest') group_name = "juxta_spike_values" if group_name in beer_file.keys(): beer_file.pop(group_name) beer_file.create_group(group_name) key = "{}/elec_0".format(group_name) beer_file.create_dataset(key, data=juxta_spike_values) beer_file.close() return
def extract_chunk_spikes(gidx, extra_thresh, valley=True): """Detect spikes from a chunk of the extracellular traces""" loc_chunk, t_offset = data_file.get_data(gidx, chunk_size, nodes=nodes) loc_shape = len(loc_chunk) # Whiten signal. if do_spatial_whitening: loc_chunk = numpy.dot(loc_chunk, spatial_whitening) if do_temporal_whitening: loc_chunk = scipy.ndimage.filters.convolve1d(loc_chunk, temporal_whitening, axis=0, mode='constant') ##### TODO: uncomment or remove temporary zone # # For each electrode, center traces by removing the medians. # extra_medians = numpy.median(loc_chunk, axis=0) # loc_chunk = loc_chunk - extra_medians ##### end temporary zone # Preallocation for results. peak_times = N_elec * [None] peak_channels = N_elec * [None] # For each electrode. for e in xrange(N_elec): # Extract the peaks of the current chunk. threshold = extra_thresh * extra_mads[e] peak_times[e] = algo.detect_peaks(loc_chunk[:, e], threshold, valley=valley, mpd=dist_peaks) peak_channels[e] = e * numpy.ones(peak_times[e].size, dtype='int') peak_values = loc_chunk[peak_times[e], e] if valley: peak_indices = numpy.where(-10.0 * threshold <= peak_values)[0] else: peak_indices = numpy.where(peak_values <= +10.0 * threshold)[0] peak_times[e] = peak_times[e][peak_indices] peak_channels[e] = peak_channels[e][peak_indices] peak_times = numpy.concatenate(peak_times) peak_channels = numpy.concatenate(peak_channels) # Remove the useless borders. if alignment: loc_borders = (2 * template_shift, loc_shape - 2 * template_shift) else: loc_borders = (template_shift, loc_shape - template_shift) peak_flags = (loc_borders[0] <= peak_times) & (peak_times < loc_borders[1]) peak_times = numpy.compress(peak_flags, peak_times) peak_channels = numpy.compress(peak_flags, peak_channels) # Filter unique peak times. loc_peak_times = numpy.unique(peak_times) ##### TODO: remove debug zone # if gidx < 1: # numpy.save("tmp/loc_peak_times_{}_{}_.npy".format(gidx, int(extra_thresh)), loc_peak_times) ##### end debug zone n_times = len(loc_peak_times) loc_peak_flags = numpy.zeros(n_times, dtype='bool') loc_peak_elecs = numpy.zeros(n_times, dtype='int') loc_peak_values = numpy.zeros(n_times, dtype='float') if 0 < len(loc_peak_times): diff_times = loc_peak_times[-1] - loc_peak_times[0] all_times = numpy.zeros((N_elec, diff_times + 1), dtype='bool') min_times = numpy.maximum( loc_peak_times - loc_peak_times[0] - safety_time, 0) max_times = numpy.minimum( loc_peak_times - loc_peak_times[0] + safety_time + 1, diff_times) # Shuffle peaks. ##### TODO: clean temporary zone # argmax_peak = numpy.random.permutation(numpy.arange(n_times)) if valley: for i, loc_peak_time in enumerate(loc_peak_times): loc_peak_values[i] = numpy.amin( loc_chunk[loc_peak_time, :]) argmax_peak = numpy.argsort(loc_peak_values) else: for i, loc_peak_time in enumerate(loc_peak_times): loc_peak_values[i] = numpy.amax( loc_chunk[loc_peak_time, :]) argmax_peak = numpy.argsort(loc_peak_values) argmes_peak = argmax_peak[::-1] ##### end temporary zone all_indices = loc_peak_times[argmax_peak] # Select peaks with spatio-temporal masks. for peak_index, peak_time in zip(argmax_peak, all_indices): # Select electrode showing lowest amplitude. if valley: elec = numpy.argmin(loc_chunk[peak_time, :]) else: elec = numpy.argmax(loc_chunk[peak_time, :]) _, neighs = get_neighbors(params, chan=elec) if safety_space: mslice = all_times[ neighs, min_times[peak_index]:max_times[peak_index]] else: mslice = all_times[ elec, min_times[peak_index]:max_times[peak_index]] is_local_min = (elec in peak_channels[peak_times == peak_time]) if is_local_min and not mslice.any(): loc_peak_flags[peak_index] = True loc_peak_elecs[peak_index] = elec if valley: loc_peak_values[peak_index] = -loc_chunk[peak_time, elec] else: loc_peak_values[peak_index] = loc_chunk[peak_time, elec] if safety_space: all_times[ neighs, min_times[peak_index]:max_times[peak_index]] = True # all_times[elec, min_times[peak_index]:max_times[peak_index]] = True else: all_times[ elec, min_times[peak_index]:max_times[peak_index]] = True loc_peak_times = numpy.compress(loc_peak_flags, loc_peak_times) loc_peak_elecs = numpy.compress(loc_peak_flags, loc_peak_elecs) loc_peak_values = numpy.compress(loc_peak_flags, loc_peak_values) ##### TODO: remove debug zone # if gidx < 1: # numpy.save("tmp/loc_peak_times_{}_{}.npy".format(gidx, int(extra_thresh)), loc_peak_times) # numpy.save("tmp/loc_peak_elecs_{}_{}.npy".format(gidx, int(extra_thresh)), loc_peak_elecs) # numpy.save("tmp/loc_peak_values_{}_{}.npy".format(gidx, int(extra_thresh)), loc_peak_values) # numpy.save("tmp/loc_chunk_{}_{}.npy".format(gidx, int(extra_thresh)), loc_chunk) ##### end debug zone return loc_peak_times + t_offset, loc_peak_elecs, loc_peak_values
def main(filename, params, nb_cpu, nb_gpu, use_gpu): try: SHARED_MEMORY = True MPI.Win.Allocate_shared(1, 1, MPI.INFO_NULL, MPI.COMM_SELF).Free() except NotImplementedError: SHARED_MEMORY = False ################################################################# sampling_rate = params.getint('data', 'sampling_rate') N_e = params.getint('data', 'N_e') N_t = params.getint('data', 'N_t') N_total = params.getint('data', 'N_total') template_shift = params.getint('data', 'template_shift') file_out = params.get('data', 'file_out') file_out_suff = params.get('data', 'file_out_suff') sign_peaks = params.get('detection', 'peaks') matched_filter = params.getboolean('detection', 'matched-filter') spike_thresh = params.getfloat('detection', 'spike_thresh') do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = int(params.getfloat('fitting', 'chunk') * sampling_rate) gpu_only = params.getboolean('fitting', 'gpu_only') nodes, edges = io.get_nodes_and_edges(params) tmp_limits = params.get('fitting', 'amp_limits').replace('(', '').replace(')', '').split(',') tmp_limits = map(float, tmp_limits) amp_auto = params.getboolean('fitting', 'amp_auto') space_explo = params.getfloat('fitting', 'space_explo') nb_chances = params.getint('fitting', 'nb_chances') max_chunk = params.getfloat('fitting', 'max_chunk') inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.argsort(nodes) ################################################################# if use_gpu: import cudamat as cmt ## Need to properly handle multi GPU per MPI nodes? if nb_gpu > nb_cpu: gpu_id = int(comm.rank // nb_cpu) else: gpu_id = 0 cmt.cuda_set_device(gpu_id) cmt.init() cmt.cuda_sync_threads() if SHARED_MEMORY: templates = io.load_data_memshared(params, comm, 'templates', normalize=True, transpose=True) N_tm, x = templates.shape else: templates = io.load_data(params, 'templates') x, N_tm = templates.shape N_e = params.getint('data', 'N_e') N_t = params.getint('data', 'N_t') template_shift = int((N_t - 1) // 2) temp_2_shift = 2 * template_shift full_gpu = use_gpu and gpu_only n_tm = N_tm // 2 n_scalar = N_e * N_t last_spikes = numpy.zeros((n_tm, 1), dtype=numpy.int32) temp_window = numpy.arange(-template_shift, template_shift + 1) if not amp_auto: amp_limits = numpy.zeros((n_tm, 2)) amp_limits[:, 0] = tmp_limits[0] amp_limits[:, 1] = tmp_limits[1] else: amp_limits = io.load_data(params, 'limits') norm_templates = io.load_data(params, 'norm-templates') if not SHARED_MEMORY: for idx in xrange(templates.shape[1]): myslice = numpy.arange(templates.indptr[idx], templates.indptr[idx + 1]) templates.data[myslice] /= norm_templates[idx] templates = templates.T if use_gpu: templates = cmt.SparseCUDAMatrix(templates) info_string = '' if matched_filter: if sign_peaks in ['negative', 'both']: waveform_neg = io.load_data(params, 'waveform') waveform_neg /= (numpy.abs(numpy.sum(waveform_neg)) * len(waveform_neg)) matched_tresholds_neg = io.load_data(params, 'matched-thresholds') if sign_peaks in ['positive', 'both']: waveform_pos = io.load_data(params, 'waveform-pos') waveform_pos /= (numpy.abs(numpy.sum(waveform_pos)) * len(waveform_pos)) matched_tresholds_pos = io.load_data(params, 'matched-thresholds-pos') if comm.rank == 0: if use_gpu: info_string = "using %d GPUs" % (comm.size) else: info_string = "using %d CPUs" % (comm.size) comm.Barrier() thresholds = io.load_data(params, 'thresholds') if SHARED_MEMORY: c_overs = io.load_data_memshared(params, comm, 'overlaps', nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) c_overlap = io.get_overlaps(comm, params, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] else: c_overlap = io.get_overlaps(comm, params, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_x = c_overlap.get('over_x')[:] over_y = c_overlap.get('over_y')[:] over_data = c_overlap.get('over_data')[:] over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] c_overlap.close() # To be faster, we rearrange the overlaps into a dictionnary. This has a cost: twice the memory usage for # a short period of time c_overs = {} overlaps = scipy.sparse.csr_matrix( (over_data, (over_x, over_y)), shape=(over_shape[0], over_shape[1])) del over_x, over_y, over_data for i in xrange(N_over): c_overs[i] = overlaps[i * N_over:(i + 1) * N_over] del overlaps comm.Barrier() if comm.rank == 0: io.print_and_log([ "Here comes the SpyKING CIRCUS %s and %d templates..." % (info_string, n_tm) ], 'default', params) io.purge(file_out_suff, '.data') if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') if full_gpu: try: # If memory on the GPU is large enough, we load the overlaps onto it for i in xrange(N_over): c_overs[i] = cmt.SparseCUDAMatrix(c_overs[i]) except Exception: if comm.rank == 0: io.print_and_log([ "Not enough memory on GPUs: GPUs are used for projection only" ], 'info', params) for i in xrange(N_over): if c_overs.has_key(i): del c_overs[i] full_gpu = False borders, nb_chunks, chunk_len, last_chunk_len = io.analyze_data( params, chunk_size) nb_chunks = int(min(nb_chunks, max_chunk)) if comm.rank == 0: pbar = get_progressbar(int(nb_chunks // comm.size)) spiketimes_file = open(file_out_suff + '.spiketimes-%d.data' % comm.rank, 'wb') amplitudes_file = open(file_out_suff + '.amplitudes-%d.data' % comm.rank, 'wb') templates_file = open(file_out_suff + '.templates-%d.data' % comm.rank, 'wb') comm.Barrier() if use_gpu and do_spatial_whitening: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) last_chunk_size = 0 for gcount, gidx in enumerate(xrange(comm.rank, nb_chunks, comm.size)): #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." ## We need to deal with the borders by taking chunks of size [0, chunck_size+template_shift] if gidx == (nb_chunks - 1): padding = (-2 * borders, 0) elif gidx == 0: padding = (0, 2 * borders) else: padding = (-2 * borders, 2 * borders) result = {'spiketimes': [], 'amplitudes': [], 'templates': []} local_chunk, local_shape = io.load_chunk(params, gidx, chunk_len, chunk_size, padding, nodes=nodes) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d(local_chunk, temporal_whitening, axis=0, mode='constant') #print "Extracting the peaks..." local_peaktimes = numpy.zeros(0, dtype=numpy.int32) if matched_filter: if sign_peaks in ['positive', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d( local_chunk, waveform_pos, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_pos[i]) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) if sign_peaks in ['negative', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d( local_chunk, waveform_neg, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_neg[i]) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) else: for i in xrange(N_e): if sign_peaks == 'negative': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=True) elif sign_peaks == 'positive': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=False) elif sign_peaks == 'both': peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False) local_peaktimes = numpy.concatenate( (local_peaktimes, peaktimes)) local_peaktimes = numpy.unique(local_peaktimes) #print "Removing the useless borders..." local_borders = (template_shift, local_shape - template_shift) idx = (local_peaktimes >= local_borders[0]) & (local_peaktimes < local_borders[1]) local_peaktimes = numpy.compress(idx, local_peaktimes) n_t = len(local_peaktimes) len_chunk = local_chunk.shape[0] all_indices = numpy.arange(n_t) if full_gpu: # all_indices = cmt.CUDAMatrix(all_indices) tmp_gpu = cmt.CUDAMatrix(local_peaktimes.reshape((1, n_t)), copy_on_host=False) if n_t > 0: #print "Computing the b (should full_gpu by putting all chunks on GPU if possible?)..." local_chunk = local_chunk.T.ravel() sub_mat = numpy.zeros((N_e * (2 * template_shift + 1), n_t), dtype=numpy.float32) if len_chunk != last_chunk_size: slice_indices = numpy.zeros(0, dtype=numpy.int32) for idx in xrange(N_e): slice_indices = numpy.concatenate( (slice_indices, len_chunk * idx + temp_window)) last_chunk_size = len_chunk for count, idx in enumerate(local_peaktimes): sub_mat[:, count] = numpy.take(local_chunk, slice_indices + idx) del local_chunk if use_gpu: sub_mat = cmt.CUDAMatrix(sub_mat, copy_on_host=False) b = cmt.sparse_dot(templates, sub_mat) else: b = templates.dot(sub_mat) del sub_mat local_offset = gidx * chunk_size + padding[0] // N_total local_bounds = (temp_2_shift, local_shape - temp_2_shift) all_spikes = local_peaktimes + local_offset penalty = numpy.ones((n_tm, n_t), dtype=numpy.float32) # Because for GPU, slicing by columns is more efficient, we need to transpose b #b = b.transpose() if use_gpu and not full_gpu: b = b.asarray() failure = numpy.zeros(n_t, dtype=numpy.int32) if full_gpu: mask = cmt.CUDAMatrix(penalty, copy_on_host=False) data = cmt.empty(mask.shape) cm_zeros = cmt.CUDAMatrix(numpy.zeros(mask.shape), copy_on_host=False) patch_gpu = b.shape[1] == 1 else: mask = penalty sub_b = b[:n_tm, :] min_time = local_peaktimes.min() max_time = local_peaktimes.max() local_len = max_time - min_time + 1 min_times = numpy.maximum( local_peaktimes - min_time - temp_2_shift, 0) max_times = numpy.minimum( local_peaktimes - min_time + temp_2_shift + 1, max_time - min_time) max_n_t = int(space_explo * (max_time - min_time + 1) // (2 * temp_2_shift + 1)) while (numpy.mean(failure) < nb_chances): if full_gpu: sub_b = b.get_row_slice(0, n_tm) sub_b.mult(mask, data) tmp_mat = data.max(0) argmax_bi = numpy.argsort(tmp_mat.asarray()[0, :])[::-1] del tmp_mat, sub_b else: data = sub_b * mask argmax_bi = numpy.argsort(numpy.max(data, 0))[::-1] while (len(argmax_bi) > 0): subset = [] indices = [] all_times = numpy.zeros(local_len, dtype=numpy.bool) for count, idx in enumerate(argmax_bi): myslice = all_times[min_times[idx]:max_times[idx]] if not myslice.any(): subset += [idx] indices += [count] all_times[min_times[idx]:max_times[idx]] = True if len(subset) > max_n_t: break subset = numpy.array(subset, dtype=numpy.int32) argmax_bi = numpy.delete(argmax_bi, indices) if full_gpu: sub_b = b.get_row_slice(0, n_tm) tmp_mat = sub_b.argmax(0) inds_t, inds_temp = subset, tmp_mat.asarray()[ 0, :][subset].astype(numpy.int32) del tmp_mat else: inds_t, inds_temp = subset, numpy.argmax( numpy.take(sub_b, subset, axis=1), 0) if full_gpu: best_amp = sub_b.asarray()[inds_temp, inds_t] / n_scalar best_amp2 = b.asarray()[inds_temp + n_tm, inds_t] / n_scalar sub_mask = numpy.ones((sub_b.shape), dtype=numpy.float32) sub_mask[inds_temp, inds_t] = 0 sub_mask = cmt.CUDAMatrix(sub_mask, copy_on_host=False) mask.mult(sub_mask) del sub_mask else: mask[inds_temp, inds_t] = 0 best_amp = sub_b[inds_temp, inds_t] / n_scalar best_amp2 = b[inds_temp + n_tm, inds_t] / n_scalar best_amp_n = best_amp / numpy.take(norm_templates, inds_temp) best_amp2_n = best_amp2 / numpy.take( norm_templates, inds_temp + n_tm) all_idx = ((best_amp_n >= amp_limits[inds_temp, 0]) & (best_amp_n <= amp_limits[inds_temp, 1])) to_keep = numpy.where(all_idx == True)[0] to_reject = numpy.where(all_idx == False)[0] ts = numpy.take(local_peaktimes, inds_t[to_keep]) good = (ts >= local_bounds[0]) & (ts < local_bounds[1]) # We reduce to only the good times that will be kept #to_keep = to_keep[good] #ts = ts[good] if len(ts) > 0: if full_gpu: tmp = cmt.CUDAMatrix(numpy.ones((len(ts), 1)), copy_on_host=False) tmp3 = cmt.CUDAMatrix(-ts.reshape((len(ts), 1)), copy_on_host=False) tmp = tmp.dot(tmp_gpu) tmp.add_col_vec(tmp3) condition = cmt.empty(tmp.shape) cmt.abs(tmp, condition).less_than(temp_2_shift + 1) condition = condition.asarray().astype(numpy.bool) tmp = tmp.asarray().astype(numpy.int32) else: tmp = numpy.dot( numpy.ones((len(ts), 1), dtype=numpy.int32), local_peaktimes.reshape((1, n_t))) tmp -= ts.reshape((len(ts), 1)) condition = numpy.abs(tmp) <= temp_2_shift for count, keep in enumerate(to_keep): idx_b = numpy.compress(condition[count, :], all_indices) ytmp = tmp[count, condition[count, :]] + temp_2_shift indices = numpy.zeros((S_over, len(ytmp)), dtype=numpy.float32) indices[ytmp, numpy.arange(len(ytmp))] = 1 if full_gpu: indices = cmt.CUDAMatrix(indices, copy_on_host=False) if patch_gpu: b_lines = b.get_col_slice(0, b.shape[0]) else: b_lines = b.get_col_slice( idx_b[0], idx_b[-1] + 1) tmp1 = cmt.sparse_dot(c_overs[inds_temp[keep]], indices, mult=-best_amp[keep]) tmp2 = cmt.sparse_dot(c_overs[inds_temp[keep] + n_tm], indices, mult=-best_amp2[keep]) b_lines.add(tmp1) b_lines.add(tmp2) del tmp1, tmp2 else: tmp1 = c_overs[inds_temp[keep]].multiply( -best_amp[keep]).dot(indices) tmp2 = c_overs[inds_temp[keep] + n_tm].multiply(-best_amp2[keep] ).dot(indices) b[:, idx_b] += tmp1 + tmp2 if good[count]: t_spike = ts[count] + local_offset result['spiketimes'] += [t_spike] result['amplitudes'] += [(best_amp_n[keep], best_amp2_n[keep])] result['templates'] += [inds_temp[keep]] myslice = numpy.take(inds_t, to_reject) failure[myslice] += 1 sub_idx = (numpy.take(failure, myslice) >= nb_chances) if full_gpu: N = numpy.sum(sub_idx) if N > 0: cu_slice = cmt.CUDAMatrix(numpy.compress( sub_idx, myslice).reshape(1, N), copy_on_host=False) mask.set_selected_columns(cu_slice, cm_zeros) del cu_slice else: mask[:, numpy.compress(sub_idx, myslice)] = 0 if full_gpu: del sub_b spikes_to_write = numpy.array(result['spiketimes'], dtype=numpy.uint32) amplitudes_to_write = numpy.array(result['amplitudes'], dtype=numpy.float32) templates_to_write = numpy.array(result['templates'], dtype=numpy.int32) spiketimes_file.write(spikes_to_write.tostring()) amplitudes_file.write(amplitudes_to_write.tostring()) templates_file.write(templates_to_write.tostring()) if full_gpu: del mask, b, cm_zeros, data if comm.rank == 0: pbar.update(gcount) spiketimes_file.flush() os.fsync(spiketimes_file.fileno()) spiketimes_file.close() amplitudes_file.flush() os.fsync(amplitudes_file.fileno()) amplitudes_file.close() templates_file.flush() os.fsync(templates_file.fileno()) templates_file.close() comm.Barrier() if comm.rank == 0: pbar.finish() if comm.rank == 0: io.collect_data(comm.size, params, erase=True)
def main(params, nb_cpu, nb_gpu, use_gpu): ################################################################# logger = init_logging(params.logfile) logger = logging.getLogger('circus.fitting') data_file = params.data_file data_file.open() N_e = params.getint('data', 'N_e') N_total = params.nb_channels N_t = params.getint('detection', 'N_t') template_shift = params.getint('detection', 'template_shift') file_out = params.get('data', 'file_out') file_out_suff = params.get('data', 'file_out_suff') sign_peaks = params.get('detection', 'peaks') matched_filter = params.getboolean('detection', 'matched-filter') spike_thresh = params.getfloat('detection', 'spike_thresh') do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = params.getint('fitting', 'chunk_size') gpu_only = params.getboolean('fitting', 'gpu_only') nodes, edges = get_nodes_and_edges(params) tmp_limits = params.get('fitting', 'amp_limits').replace('(', '').replace(')', '').split(',') tmp_limits = map(float, tmp_limits) amp_auto = params.getboolean('fitting', 'amp_auto') space_explo = params.getfloat('fitting', 'space_explo') nb_chances = params.getint('fitting', 'nb_chances') max_chunk = params.getfloat('fitting', 'max_chunk') noise_thr = params.getfloat('clustering', 'noise_thr') collect_all = params.getboolean('fitting', 'collect_all') ignore_dead_times = params.getboolean('triggers', 'ignore_times') inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.argsort(nodes) ################################################################# if use_gpu: import cudamat as cmt ## Need to properly handle multi GPU per MPI nodes? if nb_gpu > nb_cpu: gpu_id = int(comm.rank//nb_cpu) else: gpu_id = 0 cmt.cuda_set_device(gpu_id) cmt.init() cmt.cuda_sync_threads() if SHARED_MEMORY: templates = io.load_data_memshared(params, 'templates', normalize=True, transpose=True) N_tm, x = templates.shape else: templates = io.load_data(params, 'templates') x, N_tm = templates.shape temp_2_shift = 2*template_shift full_gpu = use_gpu and gpu_only n_tm = N_tm//2 n_scalar = N_e*N_t last_spikes = numpy.zeros((n_tm, 1), dtype=numpy.int32) temp_window = numpy.arange(-template_shift, template_shift+1) if not amp_auto: amp_limits = numpy.zeros((n_tm, 2)) amp_limits[:, 0] = tmp_limits[0] amp_limits[:, 1] = tmp_limits[1] else: amp_limits = io.load_data(params, 'limits') norm_templates = io.load_data(params, 'norm-templates') if not SHARED_MEMORY: for idx in xrange(templates.shape[1]): myslice = numpy.arange(templates.indptr[idx], templates.indptr[idx+1]) templates.data[myslice] /= norm_templates[idx] templates = templates.T if matched_filter: if sign_peaks in ['negative', 'both']: waveform_neg = io.load_data(params, 'waveform') waveform_neg /= (numpy.abs(numpy.sum(waveform_neg))* len(waveform_neg)) matched_tresholds_neg = io.load_data(params, 'matched-thresholds') if sign_peaks in ['positive', 'both']: waveform_pos = io.load_data(params, 'waveform-pos') waveform_pos /= (numpy.abs(numpy.sum(waveform_pos))* len(waveform_pos)) matched_tresholds_pos = io.load_data(params, 'matched-thresholds-pos') if ignore_dead_times: dead_times = numpy.loadtxt(params.get('triggers', 'dead_file')) if len(dead_times.shape) == 1: dead_times = dead_times.reshape(1, 2) dead_in_ms = params.getboolean('triggers', 'dead_in_ms') if dead_in_ms: dead_times *= numpy.int64(data_file.sampling_rate*1e-3) dead_times = dead_times.astype(numpy.int64) all_dead_times = [] for i in xrange(len(dead_times)): all_dead_times += range(dead_times[i, 0], dead_times[i, 1]) thresholds = io.load_data(params, 'thresholds') if collect_all: neighbors = {} for i in xrange(n_tm): tmp = templates[i, :].toarray().reshape(N_e, N_t) * norm_templates[i] neighbors[i] = numpy.where(numpy.sum(tmp, 1) != 0)[0] if use_gpu: templates = cmt.SparseCUDAMatrix(templates, copy_on_host=False) info_string = '' if comm.rank == 0: if use_gpu: info_string = "using %d GPUs" %(comm.size) else: info_string = "using %d CPUs" %(comm.size) comm.Barrier() c_overlap = io.get_overlaps(params, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] ## If the number of overlaps is different from templates, we need to recompute them if N_over != N_tm: if comm.rank == 0: print_and_log(['Templates have been modified, recomputing the overlaps...'], 'default', logger) c_overlap = io.get_overlaps(params, erase=True, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_shape = c_overlap.get('over_shape')[:] N_over = int(numpy.sqrt(over_shape[0])) S_over = over_shape[1] if SHARED_MEMORY: c_overs = io.load_data_memshared(params, 'overlaps', nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) else: c_overlap = io.get_overlaps(params, nb_cpu=nb_cpu, nb_gpu=nb_gpu, use_gpu=use_gpu) over_x = c_overlap.get('over_x')[:] over_y = c_overlap.get('over_y')[:] over_data = c_overlap.get('over_data')[:] over_shape = c_overlap.get('over_shape')[:] c_overlap.close() # To be faster, we rearrange the overlaps into a dictionnary. This has a cost: twice the memory usage for # a short period of time c_overs = {} overlaps = scipy.sparse.csr_matrix((over_data, (over_x, over_y)), shape=(over_shape[0], over_shape[1])) del over_x, over_y, over_data for i in xrange(N_over): c_overs[i] = overlaps[i*N_over:(i+1)*N_over] del overlaps comm.Barrier() if comm.rank == 0: print_and_log(["Here comes the SpyKING CIRCUS %s and %d templates..." %(info_string, n_tm)], 'default', logger) purge(file_out_suff, '.data') if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') if full_gpu: try: # If memory on the GPU is large enough, we load the overlaps onto it for i in xrange(N_over): c_overs[i] = cmt.SparseCUDAMatrix(c_overs[i], copy_on_host=False) except Exception: if comm.rank == 0: print_and_log(["Not enough memory on GPUs: GPUs are used for projection only"], 'info', logger) for i in xrange(N_over): if c_overs.has_key(i): del c_overs[i] full_gpu = False nb_chunks, last_chunk_len = data_file.analyze(chunk_size) processed_chunks = int(min(nb_chunks, max_chunk)) comm.Barrier() spiketimes_file = open(file_out_suff + '.spiketimes-%d.data' %comm.rank, 'wb') comm.Barrier() amplitudes_file = open(file_out_suff + '.amplitudes-%d.data' %comm.rank, 'wb') comm.Barrier() templates_file = open(file_out_suff + '.templates-%d.data' %comm.rank, 'wb') comm.Barrier() if collect_all: garbage_times_file = open(file_out_suff + '.gspiketimes-%d.data' %comm.rank, 'wb') comm.Barrier() garbage_temp_file = open(file_out_suff + '.gtemplates-%d.data' %comm.rank, 'wb') comm.Barrier() if use_gpu and do_spatial_whitening: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) last_chunk_size = 0 to_explore = xrange(comm.rank, processed_chunks, comm.size) if comm.rank == 0: to_explore = get_tqdm_progressbar(to_explore) for gcount, gidx in enumerate(to_explore): #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." ## We need to deal with the borders by taking chunks of size [0, chunck_size+template_shift] is_first = data_file.is_first_chunk(gidx, nb_chunks) is_last = data_file.is_last_chunk(gidx, nb_chunks) if is_last: padding = (-2*template_shift, 0) elif is_first: padding = (0, 2*template_shift) else: padding = (-2*template_shift, 2*template_shift) result = {'spiketimes' : [], 'amplitudes' : [], 'templates' : []} local_chunk, t_offset = data_file.get_data(gidx, chunk_size, padding, nodes=nodes) len_chunk = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d(local_chunk, temporal_whitening, axis=0, mode='constant') #print "Extracting the peaks..." if collect_all: all_found_spikes = {} for i in xrange(N_e): all_found_spikes[i] = [] local_peaktimes = numpy.zeros(0, dtype=numpy.int32) if matched_filter: if sign_peaks in ['positive', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_pos, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_pos[i]) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() if sign_peaks in ['negative', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_neg, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_neg[i]) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() else: for i in xrange(N_e): if sign_peaks == 'negative': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=True) elif sign_peaks == 'positive': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=False) elif sign_peaks == 'both': peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) if collect_all: all_found_spikes[i] += peaktimes.tolist() local_peaktimes = numpy.unique(local_peaktimes) if ignore_dead_times: local_peaktimes = numpy.array(list(set(local_peaktimes + t_offset).difference(all_dead_times)), dtype=numpy.int32) - t_offset local_peaktimes = numpy.sort(local_peaktimes) #print "Removing the useless borders..." local_borders = (template_shift, len_chunk - template_shift) idx = (local_peaktimes >= local_borders[0]) & (local_peaktimes < local_borders[1]) local_peaktimes = numpy.compress(idx, local_peaktimes) if collect_all: for i in xrange(N_e): all_found_spikes[i] = numpy.array(all_found_spikes[i], dtype=numpy.int32) if ignore_dead_times: all_found_spikes[i] = numpy.array(list(set(all_found_spikes[i] + t_offset).difference(all_dead_times)), dtype=numpy.int32) - t_offset all_found_spikes[i] = numpy.sort(all_found_spikes[i]) idx = (all_found_spikes[i] >= local_borders[0]) & (all_found_spikes[i] < local_borders[1]) all_found_spikes[i] = numpy.compress(idx, all_found_spikes[i]) n_t = len(local_peaktimes) all_indices = numpy.arange(n_t) if full_gpu: # all_indices = cmt.CUDAMatrix(all_indices) tmp_gpu = cmt.CUDAMatrix(local_peaktimes.reshape((1, n_t)), copy_on_host=False) if n_t > 0: #print "Computing the b (should full_gpu by putting all chunks on GPU if possible?)..." if collect_all: c_local_chunk = local_chunk.copy() local_chunk = local_chunk.T.ravel() sub_mat = numpy.zeros((N_e*(2*template_shift+1), n_t), dtype=numpy.float32) if len_chunk != last_chunk_size: slice_indices = numpy.zeros(0, dtype=numpy.int32) for idx in xrange(N_e): slice_indices = numpy.concatenate((slice_indices, len_chunk*idx + temp_window)) last_chunk_size = len_chunk for count, idx in enumerate(local_peaktimes): sub_mat[:, count] = numpy.take(local_chunk, slice_indices + idx) del local_chunk if use_gpu: sub_mat = cmt.CUDAMatrix(sub_mat, copy_on_host=False) b = cmt.sparse_dot(templates, sub_mat) else: b = templates.dot(sub_mat) del sub_mat local_offset = padding[0] + t_offset local_bounds = (temp_2_shift, len_chunk - temp_2_shift) all_spikes = local_peaktimes + local_offset # Because for GPU, slicing by columns is more efficient, we need to transpose b #b = b.transpose() if use_gpu and not full_gpu: b = b.asarray() failure = numpy.zeros(n_t, dtype=numpy.int32) if full_gpu: mask = numpy.zeros((2*n_tm, n_t), dtype=numpy.float32) mask[:n_tm, :] = 1 data = cmt.empty(mask.shape) patch_gpu= b.shape[1] == 1 else: mask = numpy.ones((n_tm, n_t), dtype=numpy.float32) sub_b = b[:n_tm, :] min_time = local_peaktimes.min() max_time = local_peaktimes.max() local_len = max_time - min_time + 1 min_times = numpy.maximum(local_peaktimes - min_time - temp_2_shift, 0) max_times = numpy.minimum(local_peaktimes - min_time + temp_2_shift + 1, max_time - min_time) max_n_t = int(space_explo*(max_time-min_time+1)//(2*temp_2_shift + 1)) if collect_all: c_all_times = numpy.zeros((len_chunk, N_e), dtype=numpy.bool) c_min_times = numpy.maximum(numpy.arange(len_chunk) - template_shift, 0) c_max_times = numpy.minimum(numpy.arange(len_chunk) + template_shift + 1, len_chunk) for i in xrange(N_e): c_all_times[all_found_spikes[i], i] = True while (numpy.mean(failure) < nb_chances): if full_gpu: gpu_mask = cmt.CUDAMatrix(mask, copy_on_host=False) b.mult(gpu_mask, data) tmp_mat = data.max(0) argmax_bi = numpy.argsort(tmp_mat.asarray()[0, :])[::-1] del tmp_mat else: data = sub_b * mask argmax_bi = numpy.argsort(numpy.max(data, 0))[::-1] while (len(argmax_bi) > 0): subset = [] indices = [] all_times = numpy.zeros(local_len, dtype=numpy.bool) for count, idx in enumerate(argmax_bi): myslice = all_times[min_times[idx]:max_times[idx]] if not myslice.any(): subset += [idx] indices += [count] all_times[min_times[idx]:max_times[idx]] = True if len(subset) > max_n_t: break subset = numpy.array(subset, dtype=numpy.int32) argmax_bi = numpy.delete(argmax_bi, indices) if full_gpu: b_array = b.asarray() sub_b = b_array[:n_tm, :] inds_t, inds_temp = subset, numpy.argmax(numpy.take(sub_b, subset, axis=1), 0) if full_gpu: best_amp = sub_b[inds_temp, inds_t]/n_scalar best_amp2 = b_array[inds_temp + n_tm, inds_t]/n_scalar else: best_amp = sub_b[inds_temp, inds_t]/n_scalar best_amp2 = b[inds_temp + n_tm, inds_t]/n_scalar mask[inds_temp, inds_t] = 0 best_amp_n = best_amp/numpy.take(norm_templates, inds_temp) best_amp2_n = best_amp2/numpy.take(norm_templates, inds_temp + n_tm) all_idx = ((best_amp_n >= amp_limits[inds_temp, 0]) & (best_amp_n <= amp_limits[inds_temp, 1])) to_keep = numpy.where(all_idx == True)[0] to_reject = numpy.where(all_idx == False)[0] ts = numpy.take(local_peaktimes, inds_t[to_keep]) good = (ts >= local_bounds[0]) & (ts < local_bounds[1]) # We reduce to only the good times that will be kept #to_keep = to_keep[good] #ts = ts[good] if len(ts) > 0: if full_gpu: tmp = cmt.CUDAMatrix(numpy.ones((len(ts), 1)), copy_on_host=False) tmp3 = cmt.CUDAMatrix(-ts.reshape((len(ts), 1)), copy_on_host=False) tmp = tmp.dot(tmp_gpu) tmp.add_col_vec(tmp3) condition = cmt.empty(tmp.shape) cmt.abs(tmp, condition).less_than(temp_2_shift + 1) condition = condition.asarray().astype(numpy.bool) tmp = tmp.asarray().astype(numpy.int32) else: tmp = numpy.dot(numpy.ones((len(ts), 1), dtype=numpy.int32), local_peaktimes.reshape((1, n_t))) tmp -= ts.reshape((len(ts), 1)) condition = numpy.abs(tmp) <= temp_2_shift for count, keep in enumerate(to_keep): idx_b = numpy.compress(condition[count, :], all_indices) ytmp = tmp[count, condition[count, :]] + temp_2_shift indices = numpy.zeros((S_over, len(ytmp)), dtype=numpy.float32) indices[ytmp, numpy.arange(len(ytmp))] = 1 if full_gpu: indices = cmt.CUDAMatrix(indices, copy_on_host=False) if patch_gpu: b_lines = b.get_col_slice(0, b.shape[0]) else: b_lines = b.get_col_slice(idx_b[0], idx_b[-1]+1) tmp1 = cmt.sparse_dot(c_overs[inds_temp[keep]], indices, mult=-best_amp[keep]) tmp2 = cmt.sparse_dot(c_overs[inds_temp[keep] + n_tm], indices, mult=-best_amp2[keep]) b_lines.add(tmp1.add(tmp2)) del tmp1, tmp2 else: tmp1 = c_overs[inds_temp[keep]].multiply(-best_amp[keep]).dot(indices) tmp2 = c_overs[inds_temp[keep] + n_tm].multiply(-best_amp2[keep]).dot(indices) b[:, idx_b] += tmp1 + tmp2 if good[count]: t_spike = ts[count] + local_offset result['spiketimes'] += [t_spike] result['amplitudes'] += [(best_amp_n[keep], best_amp2_n[keep])] result['templates'] += [inds_temp[keep]] myslice = numpy.take(inds_t, to_reject) failure[myslice] += 1 sub_idx = (numpy.take(failure, myslice) >= nb_chances) mask[:, numpy.compress(sub_idx, myslice)] = 0 spikes_to_write = numpy.array(result['spiketimes'], dtype=numpy.uint32) amplitudes_to_write = numpy.array(result['amplitudes'], dtype=numpy.float32) templates_to_write = numpy.array(result['templates'], dtype=numpy.int32) spiketimes_file.write(spikes_to_write.tostring()) amplitudes_file.write(amplitudes_to_write.tostring()) templates_file.write(templates_to_write.tostring()) if collect_all: for temp, spike in zip(templates_to_write, spikes_to_write - local_offset): c_all_times[c_min_times[spike]:c_max_times[spike], neighbors[temp]] = False gspikes = numpy.where(numpy.sum(c_all_times, 1) > 0)[0] c_all_times = numpy.take(c_all_times, gspikes, axis=0) c_local_chunk = numpy.take(c_local_chunk, gspikes, axis=0) * c_all_times if sign_peaks == 'negative': bestlecs = numpy.argmin(c_local_chunk, 1) if matched_filter: threshs = -matched_tresholds_neg[bestlecs] else: threshs = -thresholds[bestlecs] idx = numpy.where(numpy.min(c_local_chunk, 1) < threshs)[0] elif sign_peaks == 'positive': bestlecs = numpy.argmax(c_local_chunk, 1) if matched_filter: threshs = matched_tresholds_pos[bestlecs] else: threshs = thresholds[bestlecs] idx = numpy.where(numpy.max(c_local_chunk, 1) > threshs)[0] elif sign_peaks == 'both': c_local_chunk = numpy.abs(c_local_chunk) bestlecs = numpy.argmax(c_local_chunk, 1) if matched_filter: threshs = numpy.minimum(matched_tresholds_neg[bestlecs], matched_tresholds_pos[bestlecs]) else: threshs = thresholds[bestlecs] idx = numpy.where(numpy.max(c_local_chunk, 1) > threshs)[0] gspikes = numpy.take(gspikes, idx) bestlecs = numpy.take(bestlecs, idx) gspikes_to_write = numpy.array(gspikes + local_offset, dtype=numpy.uint32) gtemplates_to_write = numpy.array(bestlecs, dtype=numpy.int32) garbage_times_file.write(gspikes_to_write.tostring()) garbage_temp_file.write(gtemplates_to_write.tostring()) if full_gpu: del gpu_mask, b, data spiketimes_file.flush() os.fsync(spiketimes_file.fileno()) spiketimes_file.close() amplitudes_file.flush() os.fsync(amplitudes_file.fileno()) amplitudes_file.close() templates_file.flush() os.fsync(templates_file.fileno()) templates_file.close() if collect_all: garbage_temp_file.flush() os.fsync(garbage_temp_file.fileno()) garbage_temp_file.close() garbage_times_file.flush() os.fsync(garbage_times_file.fileno()) garbage_times_file.close() comm.Barrier() if comm.rank == 0: io.collect_data(comm.size, params, erase=True) data_file.close()
def main(params, nb_cpu, nb_gpu, use_gpu): ################################################################# #params = detect_memory(params) logger = init_logging(params.logfile) SHARED_MEMORY = get_shared_memory_flag(params) logger = logging.getLogger('circus.fitting') data_file = params.data_file N_e = params.getint('data', 'N_e') N_total = params.nb_channels N_t = params.getint('detection', 'N_t') template_shift = params.getint('detection', 'template_shift') file_out = params.get('data', 'file_out') file_out_suff = params.get('data', 'file_out_suff') sign_peaks = params.get('detection', 'peaks') dist_peaks = params.getint('detection', 'dist_peaks') matched_filter = params.getboolean('detection', 'matched-filter') spike_thresh = params.getfloat('detection', 'spike_thresh') do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') chunk_size = detect_memory(params) gpu_only = params.getboolean('fitting', 'gpu_only') nodes, edges = get_nodes_and_edges(params) tmp_limits = params.get('fitting', 'amp_limits').replace('(', '').replace(')', '').split(',') tmp_limits = map(float, tmp_limits) amp_auto = params.getboolean('fitting', 'amp_auto') nb_chances = params.getint('fitting', 'nb_chances') max_chunk = params.getfloat('fitting', 'max_chunk') noise_thr = params.getfloat('clustering', 'noise_thr') collect_all = params.getboolean('fitting', 'collect_all') ignore_dead_times = params.getboolean('triggers', 'ignore_times') inv_nodes = numpy.zeros(N_total, dtype=numpy.int32) inv_nodes[nodes] = numpy.arange(len(nodes)) data_file.open() ################################################################# if use_gpu: import cudamat as cmt ## Need to properly handle multi GPU per MPI nodes? if nb_gpu > nb_cpu: gpu_id = int(comm.rank//nb_cpu) else: gpu_id = 0 cmt.cuda_set_device(gpu_id) cmt.init() cmt.cuda_sync_threads() if matched_filter: if sign_peaks in ['negative', 'both']: waveform_neg = io.load_data(params, 'waveform') waveform_neg /= (numpy.abs(numpy.sum(waveform_neg))* len(waveform_neg)) matched_tresholds_neg = io.load_data(params, 'matched-thresholds') if sign_peaks in ['positive', 'both']: waveform_pos = io.load_data(params, 'waveform-pos') waveform_pos /= (numpy.abs(numpy.sum(waveform_pos))* len(waveform_pos)) matched_tresholds_pos = io.load_data(params, 'matched-thresholds-pos') if ignore_dead_times: all_dead_times = get_dead_times(params) thresholds = io.load_data(params, 'thresholds') comm.Barrier() if comm.rank == 0: print_and_log(["Extracting MUA activity..."], 'default', logger) purge(file_out_suff, '.data') if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') nb_chunks, last_chunk_len = data_file.analyze(chunk_size) processed_chunks = int(min(nb_chunks, max_chunk)) comm.Barrier() spiketimes_file = open(file_out_suff + '.mua-%d.data' %comm.rank, 'wb') comm.Barrier() electrodes_file = open(file_out_suff + '.elec-%d.data' %comm.rank, 'wb') comm.Barrier() amplitudes_file = open(file_out_suff + '.amp-%d.data' %comm.rank, 'wb') comm.Barrier() if use_gpu and do_spatial_whitening: spatial_whitening = cmt.CUDAMatrix(spatial_whitening, copy_on_host=False) to_explore = xrange(comm.rank, processed_chunks, comm.size) if comm.rank == 0: to_explore = get_tqdm_progressbar(to_explore) for gcount, gidx in enumerate(to_explore): #print "Node", comm.rank, "is analyzing chunk", gidx, "/", nb_chunks, " ..." ## We need to deal with the borders by taking chunks of size [0, chunck_size+template_shift] is_first = data_file.is_first_chunk(gidx, nb_chunks) is_last = data_file.is_last_chunk(gidx, nb_chunks) if is_last: padding = (-dist_peaks, 0) elif is_first: padding = (0, dist_peaks) else: padding = (-dist_peaks, dist_peaks) result = {'spiketimes' : [], 'amplitudes' : [], 'templates' : []} local_chunk, t_offset = data_file.get_data(gidx, chunk_size, padding, nodes=nodes) len_chunk = len(local_chunk) if do_spatial_whitening: if use_gpu: local_chunk = cmt.CUDAMatrix(local_chunk, copy_on_host=False) local_chunk = local_chunk.dot(spatial_whitening).asarray() else: local_chunk = numpy.dot(local_chunk, spatial_whitening) if do_temporal_whitening: local_chunk = scipy.ndimage.filters.convolve1d(local_chunk, temporal_whitening, axis=0, mode='constant') #print "Extracting the peaks..." local_peaktimes = numpy.zeros(0, dtype=numpy.uint32) local_elecs = numpy.zeros(0, dtype=numpy.uint32) local_amps = numpy.zeros(0, dtype=numpy.float32) if matched_filter: if sign_peaks in ['positive', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_pos, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_pos[i], mpd=dist_peaks) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) local_elecs = numpy.concatenate((local_elecs, i*numpy.ones(len(peaktimes), dtype='uint32'))) local_amps = numpy.concatenate((local_amps, filter_chunk[peaktimes, i])) if sign_peaks in ['negative', 'both']: filter_chunk = scipy.ndimage.filters.convolve1d(local_chunk, waveform_neg, axis=0, mode='constant') for i in xrange(N_e): peaktimes = algo.detect_peaks(filter_chunk[:, i], matched_tresholds_neg[i], mpd=dist_peaks) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) local_elecs = numpy.concatenate((local_elecs, i*numpy.ones(len(peaktimes), dtype='uint32'))) local_amps = numpy.concatenate((local_amps, filter_chunk[peaktimes, i])) else: for i in xrange(N_e): if sign_peaks == 'negative': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=True, mpd=dist_peaks) elif sign_peaks == 'positive': peaktimes = algo.detect_peaks(local_chunk[:, i], thresholds[i], valley=False, mpd=dist_peaks) elif sign_peaks == 'both': peaktimes = algo.detect_peaks(numpy.abs(local_chunk[:, i]), thresholds[i], valley=False, mpd=dist_peaks) local_peaktimes = numpy.concatenate((local_peaktimes, peaktimes)) local_elecs = numpy.concatenate((local_elecs, i*numpy.ones(len(peaktimes), dtype='uint32'))) local_amps = numpy.concatenate((local_amps, local_chunk[peaktimes, i])) g_offset = t_offset + padding[0] if ignore_dead_times: dead_indices = numpy.searchsorted(all_dead_times, [t_offset, t_offset + chunk_size]) if dead_indices[0] != dead_indices[1]: local_peaktimes = numpy.array(list(set(local_peaktimes + g_offset).difference(all_dead_times[dead_indices[0]:dead_indices[1]])), dtype=numpy.uint32) - g_offset local_peaktimes = numpy.sort(local_peaktimes) #print "Removing the useless borders..." local_borders = (0, len_chunk - dist_peaks) idx = (local_peaktimes >= local_borders[0]) & (local_peaktimes < local_borders[1]) local_peaktimes = numpy.compress(idx, local_peaktimes) + g_offset local_elecs = numpy.compress(idx, local_elecs) local_amps = numpy.compress(idx, local_amps) spiketimes_file.write(local_peaktimes.tostring()) electrodes_file.write(local_elecs.tostring()) amplitudes_file.write(local_amps.tostring()) sys.stderr.flush() spiketimes_file.flush() os.fsync(spiketimes_file.fileno()) spiketimes_file.close() electrodes_file.flush() os.fsync(electrodes_file.fileno()) electrodes_file.close() amplitudes_file.flush() os.fsync(amplitudes_file.fileno()) amplitudes_file.close() comm.Barrier() if comm.rank == 0: io.collect_mua(comm.size, params, erase=True) data_file.close()
def plot_extracted_extra_spikes(loc_all_chunks, data_len, mpi_input, data_dtype, chunk_len, chunk_size, N_total, nodes, extra_means, extra_stds, k, params, safety_space, safety_time): """Temporary function to see if the computed thresholds for a given dataset are valid""" count = 0 gidx = loc_all_chunks[0] loc_chunk, loc_shape = io.load_chunk(params, gidx, chunk_len, chunk_size, nodes=nodes) sampling_rate = params.getint('data', 'sampling_rate') dist_peaks = params.getint('data', 'dist_peaks') skip_artefact = params.getboolean('data', 'skip_artefact') template_shift = params.getint('data', 'template_shift') alignment = params.getboolean('detection', 'alignment') nodes, _ = io.get_nodes_and_edges(params) N_elec = nodes.size do_temporal_whitening = params.getboolean('whitening', 'temporal') do_spatial_whitening = params.getboolean('whitening', 'spatial') if do_spatial_whitening: spatial_whitening = io.load_data(params, 'spatial_whitening') if do_temporal_whitening: temporal_whitening = io.load_data(params, 'temporal_whitening') loc_chunk = numpy.zeros(data_len, dtype=data_dtype) mpi_input.Read_at(gidx * chunk_len, loc_chunk) loc_shape = chunk_size loc_chunk = loc_chunk.reshape(loc_shape, N_total) # Consider only the valid channels. loc_chunk = loc_chunk[:, nodes] # extra_means_ = extra_means[nodes] # extra_stds_ = extra_stds[nodes] extra_means_ = extra_means extra_stds_ = extra_stds # Whiten signal. if do_spatial_whitening: loc_chunk = numpy.dot(loc_chunk, spatial_whitening) if do_temporal_whitening: loc_chunk = scipy.ndimage.filters.convolve1d(loc_chunk, temporal_whitening, axis=0, mode='constant') # Preallocation for results. peak_times = N_elec * [None] peak_channels = N_elec * [None] # For each electrode. for e in xrange(0, N_elec): # Extract the peaks of the current chunk. threshold = k * extra_stds_[e] peak_times[e] = algo.detect_peaks(loc_chunk[:, e], threshold, valley=True, mpd=dist_peaks) peak_channels[e] = e * numpy.ones(peak_times[e].size, dtype='int') if skip_artefact: # Remove strong artifacts. peak_values = loc_chunk[peak_times[e], e] peak_indices = numpy.where(-10.0 * threshold <= peak_values)[0] peak_times[e] = peak_times[e][peak_indices] peak_times = numpy.concatenate(peak_times) peak_channels = numpy.concatenate(peak_channels) # Remove the useless borders. if alignment: loc_borders = (2 * template_shift, loc_shape - 2 * template_shift) else: loc_borders = (template_shift, loc_shape - template_shift) peak_flags = (loc_borders[0] <= peak_times) & (peak_times < loc_borders[1]) peak_times = peak_times[peak_flags] peak_channels = peak_channels[peak_flags] # Filter unique peak times. loc_peak_times = numpy.unique(peak_times) n_times = len(loc_peak_times) loc_peak_flags = numpy.zeros(n_times, dtype='bool') loc_peak_elecs = numpy.zeros(n_times, dtype='int') if 0 < len(loc_peak_times): diff_times = loc_peak_times[-1] - loc_peak_times[0] all_times = numpy.zeros((N_elec, diff_times + 1), dtype='bool') min_times = numpy.maximum( loc_peak_times - loc_peak_times[0] - safety_time, 0) max_times = numpy.minimum( loc_peak_times - loc_peak_times[0] + safety_time + 1, diff_times) ##### TODO: remove temporary zone numpy.random.seed(42) ##### end temporary zone argmax_peak = numpy.random.permutation(numpy.arange(n_times)) all_indices = loc_peak_times[argmax_peak] # Select peaks with spatio-temporal masks. for peak_index, peak_time in zip(argmax_peak, all_indices): # Select electrode showing lowest amplitude. elec = numpy.argmin(loc_chunk[peak_time, :]) _, neighs = get_neighbors(params, chan=elec) if safety_space: mslice = all_times[neighs, min_times[peak_index]:max_times[peak_index]] else: mslice = all_times[elec, min_times[peak_index]:max_times[peak_index]] is_local_min = (elec in peak_channels[peak_times == peak_time]) if is_local_min and not mslice.any(): loc_peak_flags[peak_index] = True loc_peak_elecs[peak_index] = elec if safety_space: all_times[ neighs, min_times[peak_index]:max_times[peak_index]] = True else: all_times[ elec, min_times[peak_index]:max_times[peak_index]] = True loc_peak_times = numpy.compress(loc_peak_flags, loc_peak_times) loc_peak_elecs = numpy.compress(loc_peak_flags, loc_peak_elecs) time = loc_peak_times channel = loc_peak_elecs pos = numpy.random.rand(time.size) - 0.5 fig = plt.figure() ax = fig.gca() # For each first chunk plot one channel per figure. for j in xrange(0, loc_chunk.shape[1]): ax = fig.gca() ax.plot(loc_chunk[:, j]) ax.plot([0, loc_chunk.shape[0] - 1], 2 * [extra_means_[j]], 'k--') ax.plot([0, loc_chunk.shape[0] - 1], 2 * [extra_means_[j] + extra_stds_[j]], 'k--') ax.plot([0, loc_chunk.shape[0] - 1], 2 * [extra_means_[j] - extra_stds_[j]], 'k--') ax.plot([0, loc_chunk.shape[0] - 1], 2 * [extra_means_[j] + k * extra_stds_[j]], 'k--') ax.plot([0, loc_chunk.shape[0] - 1], 2 * [extra_means_[j] - k * extra_stds_[j]], 'k--') idx, = numpy.where(channel == j) # y = + 250.0 * numpy.ones(idx.size) + 100.0 * pos[idx] y = -9.5 * numpy.ones(idx.size) + 1.0 * pos[idx] ax.scatter(time[idx], y, c='r') # new_idx, = numpy.where(new_channel == j) # new_y = - 350.0 * numpy.ones(new_idx.size) + 100.0 * new_pos[new_idx] # ax.scatter(new_time[new_idx], new_y, c='g') ax.set_xlim(0, loc_chunk.shape[0] - 1) # ax.set_ylim(- 400.0, 400.0) ax.set_ylim(-10.0, 10.0) plt.savefig("/tmp/check-{}-{}.png".format(j, comm.rank)) fig.clear() return