def test_BinaryRecordingExtractor():
num_seg = 2
num_chan = 3
num_samples = 30
sampling_frequency = 10000
dtype = 'int16'
files_path = [
f'test_BinaryRecordingExtractor_{i}.raw' for i in range(num_seg)
]
for i in range(num_seg):
np.memmap(files_path[i],
dtype=dtype,
mode='w+',
shape=(num_samples, num_chan))
rec = BinaryRecordingExtractor(files_path, sampling_frequency, num_chan,
dtype)
print(rec)
files_path = [
f'test_BinaryRecordingExtractor_copied_{i}.raw' for i in range(num_seg)
]
BinaryRecordingExtractor.write_recording(rec, files_path)
assert Path('test_BinaryRecordingExtractor_copied_0.raw').is_file()
def __init__(self, file_path):
# load params file related to the given shybrid recording
assert self.installed, self.installation_mesg
assert Path(file_path).suffix in [
".yml", ".yaml"
], "The 'file_path' should be a yaml file!"
params = sbio.get_params(file_path)['data']
file_path = Path(file_path)
# create a shybrid probe object
probe = sbprb.Probe(params['probe'])
nb_channels = probe.total_nb_channels
# translate the byte ordering
byte_order = params['order']
if byte_order == 'C':
time_axis = 1
elif byte_order == 'F':
time_axis = 0
bin_file = file_path.parent / f"{file_path.stem}.bin"
# piggyback on binary data recording extractor
BinaryRecordingExtractor.__init__(self,
files_path=bin_file,
sampling_frequency=float(
params['fs']),
num_chan=nb_channels,
dtype=params['dtype'],
time_axis=time_axis)
# load probe file
probegroup = read_prb(params['probe'])
self.set_probegroup(probegroup, in_place=True)
self._kwargs = {'file_path': str(Path(file_path).absolute())}
def _setup_recording(cls, recording, output_folder, params, verbose):
source_dir = Path(__file__).parent
# alias to params
p = params
experiment_name = output_folder / 'recording'
# save prb file
prb_file = output_folder / 'probe.prb'
probegroup = recording.get_probegroup()
write_prb(prb_file, probegroup, radius=p['adjacency_radius'])
# source file
if isinstance(
recording,
BinaryRecordingExtractor) and recording._kwargs['offset'] == 0:
# no need to copy
raw_filename = str(
Path(recording._kwargs['files_path'][0]).resolve())
dtype = recording._kwargs['dtype']
else:
# save binary file (chunk by hcunk) into a new file
raw_filename = output_folder / 'recording.dat'
dtype = 'int16'
BinaryRecordingExtractor.write_recording(
recording,
files_path=[raw_filename],
dtype='int16',
total_memory=p["total_memory"],
n_jobs=p["n_jobs_bin"],
verbose=False,
progress_bar=verbose)
if p['detect_sign'] < 0:
detect_sign = 'negative'
elif p['detect_sign'] > 0:
detect_sign = 'positive'
else:
detect_sign = 'both'
# set up klusta config file
with (source_dir / 'config_default.prm').open('r') as f:
klusta_config = f.readlines()
# Note: should use format with dict approach here
klusta_config = ''.join(klusta_config).format(
experiment_name, prb_file, raw_filename,
float(recording.get_sampling_frequency()),
recording.get_num_channels(), "'{}'".format(dtype),
p['threshold_strong_std_factor'], p['threshold_weak_std_factor'],
"'" + detect_sign + "'", p['extract_s_before'],
p['extract_s_after'], p['n_features_per_channel'],
p['pca_n_waveforms_max'], p['num_starting_clusters'])
with (output_folder / 'config.prm').open('w') as f:
f.writelines(klusta_config)
def test_ChannelSliceRecording():
num_seg = 2
num_chan = 3
num_samples = 30
sampling_frequency = 10000
dtype = 'int16'
file_paths = [f'test_BinaryRecordingExtractor_{i}.raw' for i in range(num_seg)]
for i in range(num_seg):
traces = np.memmap(file_paths[i], dtype=dtype, mode='w+', shape=(num_samples, num_chan))
traces[:] = np.arange(3)[None, :]
rec = BinaryRecordingExtractor(file_paths, sampling_frequency, num_chan, dtype)
# keep original ids
rec_sliced = ChannelSliceRecording(rec, channel_ids=[0, 2])
assert np.all(rec_sliced.get_channel_ids() == [0, 2])
traces = rec_sliced.get_traces(segment_index=1)
assert traces.shape[1] == 2
traces = rec_sliced.get_traces(segment_index=1, channel_ids=[0, 2])
assert traces.shape[1] == 2
traces = rec_sliced.get_traces(segment_index=1, channel_ids=[2, 0])
assert traces.shape[1] == 2
assert np.allclose(rec_sliced.get_times(0), rec.get_times(0))
# with channel ids renaming
rec_sliced2 = ChannelSliceRecording(rec, channel_ids=[0, 2], renamed_channel_ids=[3, 4])
assert np.all(rec_sliced2.get_channel_ids() == [3, 4])
traces = rec_sliced2.get_traces(segment_index=1)
assert traces.shape[1] == 2
assert np.all(traces[:, 0] == 0)
assert np.all(traces[:, 1] == 2)
traces = rec_sliced2.get_traces(segment_index=1, channel_ids=[4, 3])
assert traces.shape[1] == 2
assert np.all(traces[:, 0] == 2)
assert np.all(traces[:, 1] == 0)
# with probe and after save()
probe = pi.generate_linear_probe(num_elec=num_chan)
probe.set_device_channel_indices(np.arange(num_chan))
rec_p = rec.set_probe(probe)
rec_sliced3 = ChannelSliceRecording(rec_p, channel_ids=[0, 2], renamed_channel_ids=[3, 4])
probe3 = rec_sliced3.get_probe()
locations3 = probe3.contact_positions
cache_folder = Path('./my_cache_folder')
folder = cache_folder / 'sliced_recording'
rec_saved = rec_sliced3.save(folder=folder, chunk_size=10, n_jobs=2)
probe = rec_saved.get_probe()
assert np.array_equal(locations3, rec_saved.get_channel_locations())
traces3 = rec_saved.get_traces(segment_index=0)
assert np.all(traces3[:, 0] == 0)
assert np.all(traces3[:, 1] == 2)
def _setup_recording(cls, recording, output_folder, params, verbose):
import tridesclous as tdc
# save prb file
probegroup = recording.get_probegroup()
prb_file = output_folder / 'probe.prb'
write_prb(prb_file, probegroup)
num_seg = recording.get_num_segments()
sr = recording.get_sampling_frequency()
# source file
if isinstance(recording, BinaryRecordingExtractor
) and recording._kwargs['time_axis'] == 0:
# no need to copy
kwargs = recording._kwargs
file_paths = kwargs['file_paths']
dtype = kwargs['dtype']
num_chan = kwargs['num_chan']
file_offset = kwargs['file_offset']
else:
if verbose:
print('Local copy of recording')
# save binary file (chunk by chunk) into a new file
num_chan = recording.get_num_channels()
dtype = recording.get_dtype().str
file_paths = [
str(output_folder / f'raw_signals_{i}.raw')
for i in range(num_seg)
]
BinaryRecordingExtractor.write_recording(
recording,
file_paths=file_paths,
dtype=dtype,
total_memory=params["total_memory"],
n_jobs=params["n_jobs_bin"],
verbose=False,
progress_bar=verbose)
file_offset = 0
# initialize source and probe file
tdc_dataio = tdc.DataIO(dirname=str(output_folder))
tdc_dataio.set_data_source(type='RawData',
filenames=file_paths,
dtype=dtype,
sample_rate=float(sr),
total_channel=int(num_chan),
offset=int(file_offset))
tdc_dataio.set_probe_file(str(prb_file))
if verbose:
print(tdc_dataio)
def test_ChannelSliceRecording():
num_seg = 2
num_chan = 3
num_samples = 30
sampling_frequency = 10000
dtype = 'int16'
files_path = [
f'test_BinaryRecordingExtractor_{i}.raw' for i in range(num_seg)
]
for i in range(num_seg):
traces = np.memmap(files_path[i],
dtype=dtype,
mode='w+',
shape=(num_samples, num_chan))
traces[:] = np.arange(3)[None, :]
rec = BinaryRecordingExtractor(files_path, sampling_frequency, num_chan,
dtype)
# keep original ids
rec_sliced = ChannelSliceRecording(rec, channel_ids=[0, 2])
assert np.all(rec_sliced.get_channel_ids() == [0, 2])
traces = rec_sliced.get_traces(segment_index=1)
assert traces.shape[1] == 2
traces = rec_sliced.get_traces(segment_index=1, channel_ids=[0, 2])
assert traces.shape[1] == 2
traces = rec_sliced.get_traces(segment_index=1, channel_ids=[2, 0])
assert traces.shape[1] == 2
# with channel ids renaming
rec_sliced2 = ChannelSliceRecording(rec,
channel_ids=[0, 2],
renamed_channel_ids=[3, 4])
assert np.all(rec_sliced2.get_channel_ids() == [3, 4])
traces = rec_sliced2.get_traces(segment_index=1)
assert traces.shape[1] == 2
assert np.all(traces[:, 0] == 0)
assert np.all(traces[:, 1] == 2)
traces = rec_sliced2.get_traces(segment_index=1, channel_ids=[4, 3])
assert traces.shape[1] == 2
assert np.all(traces[:, 0] == 2)
assert np.all(traces[:, 1] == 0)
def test_BaseRecording():
num_seg = 2
num_chan = 3
num_samples = 30
sampling_frequency = 10000
dtype = 'int16'
files_path = [f'test_base_recording_{i}.raw' for i in range(num_seg)]
for i in range(num_seg):
a = np.memmap(files_path[i],
dtype=dtype,
mode='w+',
shape=(num_samples, num_chan))
a[:] = np.random.randn(*a.shape).astype(dtype)
rec = BinaryRecordingExtractor(files_path, sampling_frequency, num_chan,
dtype)
print(rec)
assert rec.get_num_segments() == 2
assert rec.get_num_channels() == 3
assert np.all(rec.ids_to_indices([0, 1, 2]) == [0, 1, 2])
assert np.all(
rec.ids_to_indices([0, 1, 2], prefer_slice=True) == slice(0, 3, None))
# annotations / properties
rec.annotate(yep='yop')
assert rec.get_annotation('yep') == 'yop'
rec.set_property('quality', [1., 3.3, np.nan])
values = rec.get_property('quality')
assert np.all(values[:2] == [
1.,
3.3,
])
# dump/load dict
d = rec.to_dict()
rec2 = BaseExtractor.from_dict(d)
rec3 = load_extractor(d)
# dump/load json
rec.dump_to_json('test_BaseRecording.json')
rec2 = BaseExtractor.load('test_BaseRecording.json')
rec3 = load_extractor('test_BaseRecording.json')
# dump/load pickle
rec.dump_to_pickle('test_BaseRecording.pkl')
rec2 = BaseExtractor.load('test_BaseRecording.pkl')
rec3 = load_extractor('test_BaseRecording.pkl')
# cache to binary
cache_folder = Path('./my_cache_folder')
folder = cache_folder / 'simple_recording'
rec.save(format='binary', folder=folder)
rec2 = BaseExtractor.load_from_folder(folder)
assert 'quality' in rec2.get_property_keys()
# but also possible
rec3 = BaseExtractor.load('./my_cache_folder/simple_recording')
# cache to memory
rec4 = rec3.save(format='memory')
traces4 = rec4.get_traces(segment_index=0)
traces = rec.get_traces(segment_index=0)
assert np.array_equal(traces4, traces)
# cache joblib several jobs
rec.save(name='simple_recording_2', chunk_size=10, n_jobs=4)
# set/get Probe only 2 channels
probe = Probe(ndim=2)
positions = [[0., 0.], [0., 15.], [0, 30.]]
probe.set_contacts(positions=positions,
shapes='circle',
shape_params={'radius': 5})
probe.set_device_channel_indices([2, -1, 0])
probe.create_auto_shape()
rec2 = rec.set_probe(probe, group_mode='by_shank')
rec2 = rec.set_probe(probe, group_mode='by_probe')
positions2 = rec2.get_channel_locations()
assert np.array_equal(positions2, [[0, 30.], [0., 0.]])
probe2 = rec2.get_probe()
positions3 = probe2.contact_positions
assert np.array_equal(positions2, positions3)
# from probeinterface.plotting import plot_probe_group, plot_probe
# import matplotlib.pyplot as plt
# plot_probe(probe)
# plot_probe(probe2)
# plt.show()
# test return_scale
sampling_frequency = 30000
traces = np.zeros((1000, 5), dtype='int16')
rec_int16 = NumpyRecording([traces], sampling_frequency)
assert rec_int16.get_dtype() == 'int16'
print(rec_int16)
traces_int16 = rec_int16.get_traces()
assert traces_int16.dtype == 'int16'
# return_scaled raise error when no gain_to_uV/offset_to_uV properties
with pytest.raises(ValueError):
traces_float32 = rec_int16.get_traces(return_scaled=True)
rec_int16.set_property('gain_to_uV', [.195] * 5)
rec_int16.set_property('offset_to_uV', [0.] * 5)
traces_float32 = rec_int16.get_traces(return_scaled=True)
assert traces_float32.dtype == 'float32'
def _setup_recording(cls, recording, output_folder, params, verbose):
p = params
source_dir = Path(__file__).parent
config_default_location = os.path.join(source_dir, 'config_default.yaml')
with open(config_default_location) as file:
yass_params = yaml.load(file, Loader=yaml.FullLoader)
# update root folder
yass_params['data']['root_folder'] = str(output_folder.absolute())
# geometry
probe_file_txt = os.path.join(output_folder, 'geom.txt')
geom_txt = recording.get_channel_locations()
np.savetxt(probe_file_txt, geom_txt)
# params
yass_params['recordings']['sampling_rate'] = recording.get_sampling_frequency()
yass_params['recordings']['n_channels'] = recording.get_num_channels()
# save to int16 raw
input_file_path = os.path.join(output_folder, 'data.bin')
dtype = 'int16' # HARD CODE THIS FOR YASS
input_file_path = output_folder / 'data.bin'
BinaryRecordingExtractor.write_recording(recording, files_path=[input_file_path],
dtype=dtype, verbose=False,
total_memory=p["total_memory"], n_jobs=p["n_jobs_bin"])
retrain = False
if params['neural_nets_path'] is None:
params['neural_nets_path'] = str(output_folder / 'tmp' / 'nn_train')
retrain = True
# MERGE yass_params with self.params that could be changed by the user
merge_params = merge_params_dict(yass_params, params)
# to yaml
fname_config = output_folder / 'config.yaml'
with open(fname_config, 'w') as file:
documents = yaml.dump(merge_params, file)
# RunNN training on exisiting
neural_nets_path = p['neural_nets_path']
if retrain:
# retrain NNs
YassSorter.train(recording, output_folder, verbose)
# update NN folder location
neural_nets_path = output_folder / 'tmp' / 'nn_train'
else:
# load previous NNs
if verbose:
print("USING PREVIOUSLY TRAINED NNs FROM THIS LOCATION: ", params['neural_nets_path'])
# use previuosly trained NN folder location
neural_nets_path = Path(params['neural_nets_path'])
merge_params['neuralnetwork']['denoise']['filename'] = str(neural_nets_path.absolute() / 'denoise.pt')
merge_params['neuralnetwork']['detect']['filename'] = str(neural_nets_path.absolute() / 'detect.pt')
from pprint import pprint
pprint(merge_params)
exit()
# to yaml again (for NNs update)
fname_config = output_folder / 'config.yaml'
with open(fname_config, 'w') as file:
documents = yaml.dump(merge_params, file)
def test_BaseRecording():
num_seg = 2
num_chan = 3
num_samples = 30
sampling_frequency = 10000
dtype = 'int16'
file_paths = [f'test_base_recording_{i}.raw' for i in range(num_seg)]
for i in range(num_seg):
a = np.memmap(file_paths[i],
dtype=dtype,
mode='w+',
shape=(num_samples, num_chan))
a[:] = np.random.randn(*a.shape).astype(dtype)
rec = BinaryRecordingExtractor(file_paths, sampling_frequency, num_chan,
dtype)
assert rec.get_num_segments() == 2
assert rec.get_num_channels() == 3
assert np.all(rec.ids_to_indices([0, 1, 2]) == [0, 1, 2])
assert np.all(
rec.ids_to_indices([0, 1, 2], prefer_slice=True) == slice(0, 3, None))
# annotations / properties
rec.annotate(yep='yop')
assert rec.get_annotation('yep') == 'yop'
rec.set_channel_groups([0, 0, 1])
rec.set_property('quality', [1., 3.3, np.nan])
values = rec.get_property('quality')
assert np.all(values[:2] == [
1.,
3.3,
])
# missing property
rec.set_property('string_property', ["ciao", "bello"], ids=[0, 1])
values = rec.get_property('string_property')
assert values[2] == ""
# setting an different type raises an error
assert_raises(Exception,
rec.set_property,
key='string_property_nan',
values=["ciao", "bello"],
ids=[0, 1],
missing_value=np.nan)
# int properties without missing values raise an error
assert_raises(Exception,
rec.set_property,
key='int_property',
values=[5, 6],
ids=[1, 2])
rec.set_property('int_property', [5, 6], ids=[1, 2], missing_value=200)
values = rec.get_property('int_property')
assert values.dtype.kind == "i"
times0 = rec.get_times(segment_index=0)
# dump/load dict
d = rec.to_dict()
rec2 = BaseExtractor.from_dict(d)
rec3 = load_extractor(d)
# dump/load json
rec.dump_to_json('test_BaseRecording.json')
rec2 = BaseExtractor.load('test_BaseRecording.json')
rec3 = load_extractor('test_BaseRecording.json')
# dump/load pickle
rec.dump_to_pickle('test_BaseRecording.pkl')
rec2 = BaseExtractor.load('test_BaseRecording.pkl')
rec3 = load_extractor('test_BaseRecording.pkl')
# dump/load dict - relative
d = rec.to_dict(relative_to=".")
rec2 = BaseExtractor.from_dict(d, base_folder=".")
rec3 = load_extractor(d, base_folder=".")
# dump/load json
rec.dump_to_json('test_BaseRecording_rel.json', relative_to=".")
rec2 = BaseExtractor.load('test_BaseRecording_rel.json', base_folder=".")
rec3 = load_extractor('test_BaseRecording_rel.json', base_folder=".")
# cache to binary
cache_folder = Path('./my_cache_folder')
folder = cache_folder / 'simple_recording'
rec.save(format='binary', folder=folder)
rec2 = BaseExtractor.load_from_folder(folder)
assert 'quality' in rec2.get_property_keys()
values = rec2.get_property('quality')
assert values[0] == 1.
assert values[1] == 3.3
assert np.isnan(values[2])
groups = rec2.get_channel_groups()
assert np.array_equal(groups, [0, 0, 1])
# but also possible
rec3 = BaseExtractor.load('./my_cache_folder/simple_recording')
# cache to memory
rec4 = rec3.save(format='memory')
traces4 = rec4.get_traces(segment_index=0)
traces = rec.get_traces(segment_index=0)
assert np.array_equal(traces4, traces)
# cache joblib several jobs
folder = cache_folder / 'simple_recording2'
rec2 = rec.save(folder=folder, chunk_size=10, n_jobs=4)
traces2 = rec2.get_traces(segment_index=0)
# set/get Probe only 2 channels
probe = Probe(ndim=2)
positions = [[0., 0.], [0., 15.], [0, 30.]]
probe.set_contacts(positions=positions,
shapes='circle',
shape_params={'radius': 5})
probe.set_device_channel_indices([2, -1, 0])
probe.create_auto_shape()
rec_p = rec.set_probe(probe, group_mode='by_shank')
rec_p = rec.set_probe(probe, group_mode='by_probe')
positions2 = rec_p.get_channel_locations()
assert np.array_equal(positions2, [[0, 30.], [0., 0.]])
probe2 = rec_p.get_probe()
positions3 = probe2.contact_positions
assert np.array_equal(positions2, positions3)
assert np.array_equal(probe2.device_channel_indices, [0, 1])
# test save with probe
folder = cache_folder / 'simple_recording3'
rec2 = rec_p.save(folder=folder, chunk_size=10, n_jobs=2)
rec2 = load_extractor(folder)
probe2 = rec2.get_probe()
assert np.array_equal(probe2.contact_positions, [[0, 30.], [0., 0.]])
positions2 = rec_p.get_channel_locations()
assert np.array_equal(positions2, [[0, 30.], [0., 0.]])
traces2 = rec2.get_traces(segment_index=0)
assert np.array_equal(traces2, rec_p.get_traces(segment_index=0))
# from probeinterface.plotting import plot_probe_group, plot_probe
# import matplotlib.pyplot as plt
# plot_probe(probe)
# plot_probe(probe2)
# plt.show()
# test return_scale
sampling_frequency = 30000
traces = np.zeros((1000, 5), dtype='int16')
rec_int16 = NumpyRecording([traces], sampling_frequency)
assert rec_int16.get_dtype() == 'int16'
traces_int16 = rec_int16.get_traces()
assert traces_int16.dtype == 'int16'
# return_scaled raise error when no gain_to_uV/offset_to_uV properties
with pytest.raises(ValueError):
traces_float32 = rec_int16.get_traces(return_scaled=True)
rec_int16.set_property('gain_to_uV', [.195] * 5)
rec_int16.set_property('offset_to_uV', [0.] * 5)
traces_float32 = rec_int16.get_traces(return_scaled=True)
assert traces_float32.dtype == 'float32'
# test with t_start
rec = BinaryRecordingExtractor(file_paths,
sampling_frequency,
num_chan,
dtype,
t_starts=np.arange(num_seg) * 10.)
times1 = rec.get_times(1)
folder = cache_folder / 'recording_with_t_start'
rec2 = rec.save(folder=folder)
assert np.allclose(times1, rec2.get_times(1))
# test with time_vector
rec = BinaryRecordingExtractor(file_paths, sampling_frequency, num_chan,
dtype)
rec.set_times(np.arange(num_samples) / sampling_frequency + 30.,
segment_index=0)
rec.set_times(np.arange(num_samples) / sampling_frequency + 40.,
segment_index=1)
times1 = rec.get_times(1)
folder = cache_folder / 'recording_with_times'
rec2 = rec.save(folder=folder)
assert np.allclose(times1, rec2.get_times(1))
rec3 = load_extractor(folder)
assert np.allclose(times1, rec3.get_times(1))