import logging
import random
import string
import shutil
import tempfile
import unittest
from pathlib import Path
import numpy as np
from one.api import ONE
import one.alf.io as alfio
from one.registration import RegistrationClient
from ibllib.pipes.ephys_tasks import (EphysRegisterRaw, EphysCompressNP1, EphysCompressNP21, EphysCompressNP24,
EphysSyncRegisterRaw, EphysSyncPulses, EphysPulses, SpikeSorting)
from ibllib.tests import base
_logger = logging.getLogger('ibllib')
[docs]
class EphysTemplate(base.IntegrationTest):
[docs]
def setUp(self) -> None:
self.data_path = self.default_data_root().joinpath('ephys', 'ephys_np2', 'raw_ephys_data')
tmpdir = tempfile.TemporaryDirectory()
self.addCleanup(tmpdir.cleanup)
self.temp_dir = Path(tmpdir.name)
self.session_path = self.temp_dir.joinpath('mars', '2054-07-13', '001')
self.session_path.mkdir(parents=True)
self.probe = 'probe00'
self.one = ONE(**base.TEST_DB, mode='local')
[docs]
def copy_nidq_data(self, sync_collection='raw_ephys_data', ext='.bin', wiring=True):
target_path = self.session_path.joinpath(sync_collection)
target_path.mkdir(parents=True, exist_ok=True)
if wiring:
target_path.joinpath('_spikeglx_ephysData_g0_t0.nidq.wiring.json').touch()
meta_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.meta')
shutil.copy(meta_file, target_path.joinpath(meta_file.name))
if ext == '.bin':
bin_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.bin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name))
elif ext == '.cbin':
bin_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.cbin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name))
ch_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name))
else:
ch_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name))
[docs]
def copy_sync_files(self, sync_collection='raw_ephys_data'):
target_path = self.session_path.joinpath(sync_collection)
target_path.mkdir(parents=True, exist_ok=True)
sync_files = self.data_path.joinpath('nidq').glob('_spikeglx_sync.*')
for file in sync_files:
shutil.copy(file, target_path.joinpath(file.name))
meta_file = self.data_path.joinpath('nidq', '_spikeglx_ephysData_g0_t0.nidq.meta')
shutil.copy(meta_file, target_path.joinpath(meta_file.name))
[docs]
def copy_ap_data(self, probe='probe00', ext='.bin', meta='NP21', wiring=True):
target_path = self.session_path.joinpath('raw_ephys_data', probe)
target_path.mkdir(parents=True, exist_ok=True)
if wiring:
target_path.joinpath('_spikeglx_ephysData_g0_t0.imec0.ap.wiring.json').touch()
if ext == '.bin':
bin_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.bin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name))
elif ext == '.cbin':
bin_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.cbin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name))
ch_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name))
else:
ch_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name))
meta_file = self.data_path.joinpath('probe00', f'{meta}_meta', '_spikeglx_ephysData_g0_t0.imec0.ap.meta')
shutil.copy(meta_file, target_path.joinpath(meta_file.name))
[docs]
def copy_lf_data(self, probe='probe00', ext='.bin', meta='NP21'):
target_path = self.session_path.joinpath('raw_ephys_data', probe)
target_path.mkdir(parents=True, exist_ok=True)
if ext == '.bin':
bin_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.bin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name.replace('ap', 'lf')))
elif ext == '.cbin':
bin_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.cbin')
shutil.copy(bin_file, target_path.joinpath(bin_file.name.replace('ap', 'lf')))
ch_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name.replace('ap', 'lf')))
else:
ch_file = self.data_path.joinpath('probe00', '_spikeglx_ephysData_g0_t0.imec0.ap.ch')
shutil.copy(ch_file, target_path.joinpath(ch_file.name.replace('ap', 'lf')))
meta_file = self.data_path.joinpath('probe00', f'{meta}_meta', '_spikeglx_ephysData_g0_t0.imec0.ap.meta')
shutil.copy(meta_file, target_path.joinpath(meta_file.name.replace('ap', 'lf')))
[docs]
def check_files(self, task, ignore_ext='blablalba'):
assert len(task.signature['output_files']) != 0
for exp_files in task.signature['output_files']:
if ignore_ext not in exp_files[0]:
file = next(self.session_path.joinpath(exp_files[1]).glob(exp_files[0]), None)
assert file.exists()
assert file in task.outputs
[docs]
class TestEphysRegisterRaw(base.IntegrationTest):
[docs]
def setUp(self) -> None:
self._tempdir = tempfile.TemporaryDirectory()
self.one = ONE(**base.TEST_DB, cache_dir=self._tempdir.name, cache_rest=None)
path, self.eid = RegistrationClient(self.one).create_new_session('ZM_1743')
# Currently the task protocol of a session must contain 'ephys' in order to create an insertion!
self.one.alyx.rest('sessions', 'partial_update', id=self.eid, data={'task_protocol': 'ephys'})
# make a random session path and move the meta files into random probe names
self.session_path = Path(self._tempdir.name).joinpath(path.relative_to(self.one.cache_dir))
# make random probe names and the directories
self.probe_NP1, self.probe_NP21, self.probe_NP24 = [''.join(random.choices(string.ascii_letters, k=5)) for i in range(3)]
self.expected_probes = [self.probe_NP24 + ext for ext in ['a', 'b', 'c', 'd']] + [self.probe_NP21, self.probe_NP1]
self.expected_models = ['NP2.4'] * 4 + ['NP2.1', '3B2']
meta_path = self.data_path.joinpath('ephys', 'ephys_np2', 'raw_ephys_data', 'probe00')
shutil.copytree(meta_path.joinpath('NP1_meta'), self.session_path.joinpath(self.probe_NP1))
shutil.copytree(meta_path.joinpath('NP21_meta'), self.session_path.joinpath(self.probe_NP21))
shutil.copytree(meta_path.joinpath('NP24_meta'), self.session_path.joinpath(self.probe_NP24))
[docs]
def test_register_raw(self):
task = EphysRegisterRaw(self.session_path, one=self.one)
status = task.run()
assert status == 0
# Even if we run the task again we should get the same output
task.run()
# check all the probes have been created
for probe, model in zip(self.expected_probes, self.expected_models):
pid = self.one.alyx.rest('insertions', 'list', session=self.eid, name=probe)
assert len(pid) == 1
assert pid[0]['model'] == model
[docs]
def tearDown(self) -> None:
self._tempdir.cleanup()
self.one.alyx.rest('sessions', 'delete', id=self.eid)
for probe in self.expected_probes:
pid = self.one.alyx.rest('insertions', 'list', session=self.eid, name=probe)
if len(pid) > 0:
self.one.alyx.rest('insertions', 'delete', pid[0]['id'])
[docs]
class TestEphysSyncRegisterRaw(EphysTemplate):
[docs]
def test_compress(self):
self.copy_nidq_data(sync_collection='raw_ephys_data', ext='.bin')
task = EphysSyncRegisterRaw(self.session_path, one=self.one, sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_already_compressed(self):
self.copy_nidq_data(sync_collection='raw_ephys_data', ext='.cbin')
task = EphysSyncRegisterRaw(self.session_path, one=self.one, sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_register(self):
self.copy_nidq_data(sync_collection='raw_ephys_data', ext=None)
task = EphysSyncRegisterRaw(self.session_path, one=self.one, sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task, ignore_ext='.cbin')
[docs]
class TestEphysSyncPulses(EphysTemplate):
[docs]
def test_pulses(self):
self.copy_nidq_data(sync_collection='raw_ephys_data', ext='.cbin')
task = EphysSyncPulses(self.session_path, one=self.one, sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task)
[docs]
class TestEphysCompressNP1(EphysTemplate):
[docs]
def test_compress(self):
self.copy_ap_data(probe=self.probe, ext='.bin', meta='NP1')
self.copy_lf_data(probe=self.probe, ext='.bin', meta='NP1')
task = EphysCompressNP1(self.session_path, one=self.one, sync_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_already_compress(self):
self.copy_ap_data(probe=self.probe, ext='.cbin', meta='NP1')
self.copy_lf_data(probe=self.probe, ext='.cbin', meta='NP1')
task = EphysCompressNP1(self.session_path, one=self.one, sync_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_register(self):
self.copy_ap_data(probe=self.probe, ext=None, meta='NP1')
self.copy_lf_data(probe=self.probe, ext=None, meta='NP1')
task = EphysCompressNP1(self.session_path, one=self.one, device_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task, ignore_ext='.cbin')
[docs]
class TestEphysCompressNP21(EphysTemplate):
[docs]
def test_process(self):
self.copy_ap_data(probe=self.probe, ext='.bin', meta='NP21')
task = EphysCompressNP21(self.session_path, one=self.one, sync_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_already_processed(self):
self.copy_ap_data(probe=self.probe, ext='.cbin', meta='NP21')
task = EphysCompressNP21(self.session_path, one=self.one, sync_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_register(self):
self.copy_ap_data(probe=self.probe, ext=None, meta='NP21')
self.copy_lf_data(probe=self.probe, ext=None, meta='NP21')
task = EphysCompressNP21(self.session_path, one=self.one, device_collection='raw_ephys_data', pname=self.probe)
status = task.run()
assert status == 0
self.check_files(task, ignore_ext='.cbin')
[docs]
class TestEphysCompressNP24(EphysTemplate):
[docs]
def test_process(self):
self.copy_ap_data(probe=self.probe, ext='.bin', meta='NP24')
task = EphysCompressNP24(self.session_path, one=self.one, sync_collection='raw_ephys_data', pname=self.probe, nshanks=4)
status = task.run(delete_original=False)
assert status == 0
self.check_files(task)
# This checks the already processed
status = task.run(delete_original=False)
assert status == 0
self.check_files(task)
[docs]
class TestEphysPulses(EphysTemplate):
[docs]
def test_probe_sync_pulses(self):
self.copy_sync_files()
self.copy_ap_data(probe=self.probe, ext='.cbin', meta='NP1', wiring=False)
task = EphysPulses(self.session_path, one=self.one, pname=[self.probe], sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_probe_sync_pulses_diff_collection(self):
self.copy_sync_files(sync_collection='raw_widefield_data')
self.copy_ap_data(probe=self.probe, ext='.cbin', meta='NP1', wiring=False)
task = EphysPulses(self.session_path, one=self.one, pname=[self.probe], sync_collection='raw_widefield_data')
status = task.run()
assert status == 0
self.check_files(task)
[docs]
def test_probe_sync_NP24(self):
self.copy_sync_files(sync_collection='raw_ephys_data')
self.copy_ap_data(probe='probe00', ext='.cbin', meta='NP1', wiring=False)
self.copy_ap_data(probe='probe01', ext='.cbin', meta='NP24', wiring=False)
task = EphysCompressNP24(self.session_path, one=self.one, pname='probe01')
task.run()
shutil.rmtree(self.session_path.joinpath('raw_ephys_data', 'probe01'))
task = EphysPulses(self.session_path, pname=['probe00', 'probe01a', 'probe01b', 'probe01c', 'probe01d'],
one=self.one, sync_collection='raw_ephys_data')
status = task.run()
assert status == 0
self.check_files(task)
# Check the outputs are the same even if np1 or np2.4 processing
alfname = dict(object='sync', namespace='spikeglx')
sync_probe00 = alfio.load_object(self.session_path.joinpath('raw_ephys_data', 'probe00'), **alfname, short_keys=True)
sync_probe01a = alfio.load_object(self.session_path.joinpath('raw_ephys_data', 'probe01a'), **alfname, short_keys=True)
sync_probe01d = alfio.load_object(self.session_path.joinpath('raw_ephys_data', 'probe01d'), **alfname, short_keys=True)
assert np.array_equal(sync_probe01d['times'], sync_probe01a['times'])
assert np.array_equal(sync_probe00['polarities'], sync_probe01a['polarities'])
assert np.array_equal(sync_probe00['channels'], sync_probe01a['channels'])
# assert np.array_equal(sync_probe00['times'], sync_probe01a['times']) # TODO this doesn't pass :(
self.check_files(task)
[docs]
class TestSpikeSortingTask(unittest.TestCase):
[docs]
def test_get_ks2_version(self):
test_strings = [
'\x1b[0m15:39:37.919 [I] ibl:90 Starting Pykilosort version ibl_1.3.0, output in gnagga^[[0m\n',
]
with tempfile.TemporaryDirectory() as tmpdir:
for i, test_string in enumerate(test_strings):
filename = tmpdir + f'/kilosort_{i}.log'
with open(filename, 'w') as fid:
fid.write(test_string)
self.assertEqual('ibl_1.3.0', SpikeSorting._fetch_iblsorter_run_version(filename))