1import os
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]
13for folder_name, name0, channel_prefix in folder_names:
14 s3_folder_name = f"{s3_base_url}/{folder_name}"
15 local_folder_name = f"{folder_name}.si"
17 if not os.path.exists(local_folder_name):
18 # make parent directories if needed
19 os.makedirs(os.path.dirname(local_folder_name), exist_ok=True)
20 print(f'Downloading {s3_folder_name} to {local_folder_name}...')
21 # IMPORTANT NOTE: we may interrupt this download early because we really only need the first part.
22 download_s3_folder(s3_folder_name, local_folder_name)
24 # For now this only works with spikeinterface 0.102
25 recording = si.load(
26 local_folder_name
27 )
29 channel_ids = [
30 f'{channel_prefix}101',
31 f'{channel_prefix}102',
32 f'{channel_prefix}103',
33 f'{channel_prefix}104',
34 f'{channel_prefix}105',
35 f'{channel_prefix}106',
36 f'{channel_prefix}107',
37 f'{channel_prefix}108',
38 f'{channel_prefix}109',
39 f'{channel_prefix}110'
40 ]
42 fname = f'{name0}-ch101-110.raw.npy'
43 if not os.path.exists(fname):
44 print(f'Writing {fname}...')
45 X = recording.get_traces(channel_ids=channel_ids, start_frame=30000, end_frame=30000 + 30000 * 10)
46 print(f'X.shape = {X.shape}')
47 np.save(fname, X)