Blame prepare_datasets/aind/prepare_aind_compression_datasets.py at main - concept-collection/ephys_compression

concept-collection / ephys_compression_tests

48 lines · 1.7 KBCode BlameHistory

6864793 support multi-channelJeremy Magland 6 months ago1import os

2import spikeinterface as si

3import numpy as np

4from s3_utils import download_s3_folder

427c40e prepare datasetsJeremy Magland 6 months ago6s3_base_url = "s3://aind-benchmark-data/ephys-compression"

6864793 support multi-channelJeremy Magland 6 months ago7

6f96f2e prepare datasetsJeremy Magland 6 months ago8folder_names = [

427c40e prepare datasetsJeremy Magland 6 months ago9 # ("aind-np2/612962_2022-04-13_19-18-04_ProbeB", "aind-np2-probeB", "CH"),

b2f7bde updateJeremy Magland 6 months ago10 # ("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"),

6f96f2e prepare datasetsJeremy Magland 6 months ago12]

6864793 support multi-channelJeremy Magland 6 months ago13

427c40e prepare datasetsJeremy Magland 6 months ago14for folder_name, name0, channel_prefix in folder_names:

6f96f2e prepare datasetsJeremy Magland 6 months ago15 s3_folder_name = f"{s3_base_url}/{folder_name}"

16 local_folder_name = f"{folder_name}.si"

6864793 support multi-channelJeremy Magland 6 months ago17

6f96f2e prepare datasetsJeremy Magland 6 months ago18 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}...')

427c40e prepare datasetsJeremy Magland 6 months ago22 # IMPORTANT NOTE: we may interrupt this download early because we really only need the first part.

6f96f2e prepare datasetsJeremy Magland 6 months ago23 download_s3_folder(s3_folder_name, local_folder_name)

427c40e prepare datasetsJeremy Magland 6 months ago25 # For now this only works with spikeinterface 0.102

6f96f2e prepare datasetsJeremy Magland 6 months ago26 recording = si.load(

27 local_folder_name

28 )

30 channel_ids = [

427c40e prepare datasetsJeremy Magland 6 months ago31 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'

6f96f2e prepare datasetsJeremy Magland 6 months ago41 ]

6864793 support multi-channelJeremy Magland 6 months ago42

6f96f2e prepare datasetsJeremy Magland 6 months ago43 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)