def write_pcs(path, params, extension, N_tm, mode=0):
spikes = numpy.load(os.path.join(output_path, 'spike_times.npy'))
labels = numpy.load(os.path.join(output_path, 'spike_templates.npy'))
max_loc_channel = get_max_loc_channel(params, extension)
nb_features = params.getint('whitening', 'output_dim')
sign_peaks = params.get('detection', 'peaks')
nodes, edges = get_nodes_and_edges(params)
N_total = params.getint('data', 'N_total')
has_support = test_if_support(params, extension)
if has_support:
supports = io.load_data(params, 'supports', extension)
else:
inv_nodes = numpy.zeros(N_total, dtype=numpy.int32)
inv_nodes[nodes] = numpy.arange(len(nodes))
if export_all:
nb_templates = N_tm + N_e
else:
nb_templates = N_tm
pc_features_ind = numpy.zeros((nb_templates, max_loc_channel),
dtype=numpy.int32)
best_elec = io.load_data(params, 'electrodes', extension)
if export_all:
best_elec = numpy.concatenate((best_elec, numpy.arange(N_e)))
if has_support:
for count, support in enumerate(supports):
nb_loc = numpy.sum(support)
pc_features_ind[count, numpy.arange(nb_loc)] = numpy.where(
support == True)[0]
else:
for count, elec in enumerate(best_elec):
nb_loc = len(edges[nodes[elec]])
pc_features_ind[count, numpy.arange(nb_loc)] = inv_nodes[edges[
nodes[elec]]]
if sign_peaks in ['negative', 'both']:
basis_proj, basis_rec = io.load_data(params, 'basis')
elif sign_peaks in ['positive']:
basis_proj, basis_rec = io.load_data(params, 'basis-pos')
to_process = numpy.arange(comm.rank, nb_templates, comm.size)
all_offsets = numpy.zeros(nb_templates, dtype=numpy.int32)
for target in range(nb_templates):
if mode == 0:
all_offsets[target] = len(numpy.where(labels == target)[0])
elif mode == 1:
all_offsets[target] = min(
500, len(numpy.where(labels == target)[0]))
all_paddings = numpy.concatenate(([0], numpy.cumsum(all_offsets)))
total_pcs = numpy.sum(all_offsets)
pc_file = os.path.join(output_path, 'pc_features.npy')
pc_file_ids = os.path.join(output_path, 'pc_feature_spike_ids.npy')
from numpy.lib.format import open_memmap
if comm.rank == 0:
pc_features = open_memmap(pc_file,
shape=(total_pcs, nb_features,
max_loc_channel),
dtype=numpy.float32,
mode='w+')
if mode == 1:
pc_ids = open_memmap(pc_file_ids,
shape=(total_pcs, ),
dtype=numpy.int32,
mode='w+')
comm.Barrier()
pc_features = open_memmap(pc_file, mode='r+')
if mode == 1:
pc_ids = open_memmap(pc_file_ids, mode='r+')
to_explore = list(range(comm.rank, nb_templates, comm.size))
if comm.rank == 0:
to_explore = get_tqdm_progressbar(params, to_explore)
all_idx = numpy.zeros(0, dtype=numpy.int32)
for gcount, target in enumerate(to_explore):
count = all_paddings[target]
if mode == 1:
idx = numpy.random.permutation(
numpy.where(labels == target)[0])[:500]
pc_ids[count:count + len(idx)] = idx
elif mode == 0:
idx = numpy.where(labels == target)[0]
elec = best_elec[target]
if has_support:
if target >= len(supports):
indices = [target - N_tm]
else:
indices = numpy.where(supports[target])[0]
else:
indices = inv_nodes[edges[nodes[elec]]]
labels_i = target * numpy.ones(len(idx))
times_i = numpy.take(spikes, idx).astype(numpy.int64)
sub_data = io.get_stas(params,
times_i,
labels_i,
elec,
neighs=indices,
nodes=nodes,
auto_align=False)
pcs = numpy.dot(sub_data, basis_proj)
pcs = numpy.swapaxes(pcs, 1, 2)
if mode == 0:
pc_features[idx, :, :len(indices)] = pcs
elif mode == 1:
pc_features[count:count + len(idx), :, :len(indices)] = pcs
comm.Barrier()
if comm.rank == 0:
numpy.save(os.path.join(output_path, 'pc_feature_ind'),
pc_features_ind.astype(
numpy.uint32)) # n_templates, n_loc_chan
times = clusters_data['times'][electrode]
clusters = clusters_data['clusters'][electrode]
assert times.shape == clusters.shape, (times.shape, clusters.shape)
selection = (clusters == local_cluster)
times = times[selection]
clusters = clusters[selection]
if times.size > nb_snippets:
indices = np.random.choice(times.size, size=nb_snippets)
indices = np.sort(indices)
times = times[indices]
clusters = clusters[indices]
nodes, _ = get_nodes_and_edges(params)
inv_nodes = np.zeros(nb_electrodes, dtype=np.int)
inv_nodes[nodes] = np.arange(len(nodes))
indices = inv_nodes[nodes]
snippets = get_stas(params, times, clusters, electrode, indices, nodes=nodes)
snippets = np.transpose(snippets, axes=(0, 2, 1))
assert snippets.shape == (nb_snippets, nb_time_steps, nb_channels), \
(snippets.shape, nb_snippets, nb_time_steps, nb_channels)
# Load template.
template = load_template(args.template_id, params, extension='')
assert template.shape == (nb_time_steps, nb_channels)
# Compute the scalar products.
snippets_ = np.reshape(snippets, (nb_snippets, nb_channels * nb_time_steps))
template_ = np.reshape(template, (nb_channels * nb_time_steps, 1))
# print(snippets_.shape)
# print(template_.shape)
scalar_products_ = snippets_.dot(template_)
scalar_products = np.reshape(scalar_products_, (nb_snippets, ))
def write_templates(path, params, extension):
max_loc_channel = get_max_loc_channel(params, extension)
templates = io.load_data(params, 'templates', extension)
N_tm = templates.shape[1] // 2
nodes, edges = get_nodes_and_edges(params)
if sparse_export:
n_channels_max = 0
for t in range(N_tm):
data = numpy.sum(
numpy.sum(templates[:, t].toarray().reshape(N_e, N_t), 1)
!= 0)
if data > n_channels_max:
n_channels_max = data
else:
n_channels_max = N_e
if export_all:
to_write_sparse = numpy.zeros((N_tm + N_e, N_t, n_channels_max),
dtype=numpy.float32)
mapping_sparse = -1 * numpy.ones(
(N_tm + N_e, n_channels_max), dtype=numpy.int32)
else:
to_write_sparse = numpy.zeros((N_tm, N_t, n_channels_max),
dtype=numpy.float32)
mapping_sparse = -1 * numpy.ones(
(N_tm, n_channels_max), dtype=numpy.int32)
has_purity = test_if_purity(params, extension)
if has_purity:
purity = io.load_data(params, 'purity', extension)
f = open(os.path.join(output_path, 'cluster_purity.tsv'), 'w')
f.write('cluster_id\tpurity\n')
for i in range(N_tm):
f.write('%d\t%g\n' % (i, purity[i]))
f.close()
for t in range(N_tm):
tmp = templates[:, t].toarray().reshape(N_e, N_t).T
x, y = tmp.nonzero()
nb_loc = len(numpy.unique(y))
if sparse_export:
all_positions = numpy.zeros(y.max() + 1, dtype=numpy.int32)
all_positions[numpy.unique(y)] = numpy.arange(
nb_loc, dtype=numpy.int32)
pos = all_positions[y]
to_write_sparse[t, x, pos] = tmp[x, y]
mapping_sparse[t, numpy.arange(nb_loc)] = numpy.unique(y)
else:
pos = y
to_write_sparse[t, x, pos] = tmp[x, y]
if export_all:
garbage = io.load_data(params, 'garbage', extension)
for t in range(N_tm, N_tm + N_e):
elec = t - N_tm
spikes = garbage['gspikes'].pop('elec_%d' % elec).astype(
numpy.int64)
spikes = numpy.random.permutation(spikes)[:100]
mapping_sparse[t, 0] = t - N_tm
waveform = io.get_stas(params,
times_i=spikes,
labels_i=np.ones(len(spikes)),
src=elec,
neighs=[elec],
nodes=nodes,
mean_mode=True)
nb_loc = 1
if sparse_export:
to_write_sparse[t, :, 0] = waveform
else:
to_write_sparse[t, :, elec] = waveform
numpy.save(os.path.join(output_path, 'templates'), to_write_sparse)
if sparse_export:
numpy.save(os.path.join(output_path, 'template_ind'),
mapping_sparse)
return N_tm
def write_pcs(path, params, comm, extension, mode=0):
spikes = numpy.load(os.path.join(output_path, 'spike_times.npy'))
labels = numpy.load(os.path.join(output_path, 'spike_templates.npy'))
max_loc_channel = get_max_loc_channel(params)
nb_features = params.getint('whitening', 'output_dim')
nodes, edges = get_nodes_and_edges(params)
N_total = params.getint('data', 'N_total')
templates = load_data(params, 'templates', extension)
N_tm = templates.shape[1] // 2
pc_features_ind = numpy.zeros((N_tm, max_loc_channel),
dtype=numpy.int32)
clusters = load_data(params, 'clusters', extension)
best_elec = clusters['electrodes']
inv_nodes = numpy.zeros(N_total, dtype=numpy.int32)
inv_nodes[nodes] = numpy.argsort(nodes)
for count, elec in enumerate(best_elec):
nb_loc = len(edges[nodes[elec]])
pc_features_ind[count, numpy.arange(nb_loc)] = inv_nodes[edges[
nodes[elec]]]
basis_proj, basis_rec = load_data(params, 'basis')
to_process = numpy.arange(comm.rank, N_tm, comm.size)
all_offsets = numpy.zeros(N_tm, dtype=numpy.int32)
for target in xrange(N_tm):
if mode == 0:
all_offsets[target] = len(numpy.where(labels == target)[0])
elif mode == 1:
all_offsets[target] = min(
500, len(numpy.where(labels == target)[0]))
all_paddings = numpy.concatenate(([0], numpy.cumsum(all_offsets)))
total_pcs = numpy.sum(all_offsets)
pc_file = os.path.join(output_path, 'pc_features.npy')
pc_file_ids = os.path.join(output_path, 'pc_feature_spike_ids.npy')
from numpy.lib.format import open_memmap
if comm.rank == 0:
pc_features = open_memmap(pc_file,
shape=(total_pcs, nb_features,
max_loc_channel),
dtype=numpy.float32,
mode='w+')
if mode == 1:
pc_ids = open_memmap(pc_file_ids,
shape=(total_pcs, ),
dtype=numpy.int32,
mode='w+')
comm.Barrier()
pc_features = open_memmap(pc_file, mode='r+')
if mode == 1:
pc_ids = open_memmap(pc_file_ids, mode='r+')
if comm.rank == 0:
pbar = get_progressbar(len(to_process))
all_idx = numpy.zeros(0, dtype=numpy.int32)
for gcount, target in enumerate(to_process):
count = all_paddings[target]
if mode == 1:
idx = numpy.random.permutation(
numpy.where(labels == target)[0])[:500]
pc_ids[count:count + len(idx)] = idx
elif mode == 0:
idx = numpy.where(labels == target)[0]
elec = best_elec[target]
indices = inv_nodes[edges[nodes[elec]]]
labels_i = target * numpy.ones(len(idx))
times_i = numpy.take(spikes, idx)
sub_data = get_stas(params,
times_i,
labels_i,
elec,
neighs=indices,
nodes=nodes,
auto_align=False)
pcs = numpy.dot(sub_data, basis_proj)
pcs = numpy.swapaxes(pcs, 1, 2)
if mode == 0:
pc_features[idx, :, :len(indices)] = pcs
elif mode == 1:
pc_features[count:count + len(idx), :, :len(indices)] = pcs
if comm.rank == 0:
pbar.update(gcount)
if comm.rank == 0:
pbar.finish()
comm.Barrier()
if comm.rank == 0:
numpy.save(os.path.join(output_path, 'pc_feature_ind'),
pc_features_ind.astype(
numpy.uint32)) #n_templates, n_loc_chan
