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

concept-collection / ephys_compression_tests

47 lines · 1.6 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"),

10 ("aind-np1/625749_2022-08-03_15-15-06_ProbeA", "aind-np1-probeA", "AP")

6f96f2e prepare datasetsJeremy Magland 6 months ago11]

6864793 support multi-channelJeremy Magland 6 months ago12

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

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

15 local_folder_name = f"{folder_name}.si"

6864793 support multi-channelJeremy Magland 6 months ago16

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

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

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

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

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

26 local_folder_name

27 )

29 channel_ids = [

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

6f96f2e prepare datasetsJeremy Magland 6 months ago40 ]

6864793 support multi-channelJeremy Magland 6 months ago41

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