2import spikeinterface as si
3import numpy as np
4from s3_utils import download_s3_folder
6s3_base_url = "s3://aind-benchmark-data/ephys-compression"
8folder_names = [
9 # ("aind-np2/612962_2022-04-13_19-18-04_ProbeB", "aind-np2-probeB", "CH"),
10 # ("aind-np1/625749_2022-08-03_15-15-06_ProbeA", "aind-np1-probeA", "AP")
11 ("ibl-np1/CSHZAD026_2020-09-04_probe00", "ibl-np1-probe00", "AP"),
12]
14for folder_name, name0, channel_prefix in folder_names:
15 s3_folder_name = f"{s3_base_url}/{folder_name}"
16 local_folder_name = f"{folder_name}.si"
18 if not os.path.exists(local_folder_name):
19 # make parent directories if needed
20 os.makedirs(os.path.dirname(local_folder_name), exist_ok=True)
21 print(f'Downloading {s3_folder_name} to {local_folder_name}...')
22 # IMPORTANT NOTE: we may interrupt this download early because we really only need the first part.
23 download_s3_folder(s3_folder_name, local_folder_name)
25 # For now this only works with spikeinterface 0.102
26 recording = si.load(
27 local_folder_name
28 )
30 channel_ids = [
31 f'{channel_prefix}101',
32 f'{channel_prefix}102',
33 f'{channel_prefix}103',
34 f'{channel_prefix}104',
35 f'{channel_prefix}105',
36 f'{channel_prefix}106',
37 f'{channel_prefix}107',
38 f'{channel_prefix}108',
39 f'{channel_prefix}109',
40 f'{channel_prefix}110'
41 ]
43 fname = f'{name0}-ch101-110.raw.npy'
44 if not os.path.exists(fname):
45 print(f'Writing {fname}...')
46 X = recording.get_traces(channel_ids=channel_ids, start_frame=30000, end_frame=30000 + 30000 * 10)
47 print(f'X.shape = {X.shape}')
48 np.save(fname, X)