mea_analysis/create_single_file.py at main · cafalchio/mea_analysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
from tqdm.auto import tqdm
import numpy as np
import os
import gc
from utils import add_paths
from read_header import read_header
import logging

logging.basicConfig(filename="single.log", level=logging.DEBUG)
gc.enable()

## This script create a single binary file to be used with spyking circus.
# First it completes channels with missing data with zeros in the end.
# The inclusion of zeros could generate error on analysis that requires spike times.
#### CHANGED -> Instead of appending zeros, the channels were removed.


def create_zeros(size, good_channels, data_path, amp_names):
    """This function create zeros file to substitute files that are smaller than the majority
    In the end some of the channels will be just zeros from start to end and can be discarded,
    They are kept here to keed the size 64 channels.

    Inputs:
        size: the size of the file in int16 bytes
        good_channels: the list of channels with that size
        data_path: the path to the data
        amp_names: the list of amp_names
    """
    logging.debug("Function create_zeros")
    for filename in amp_names:
        if filename not in good_channels:
            # filesize = int(os.path.getsize(data_path+filename)/2)
            new_data = np.zeros(size)
            with open(add_paths(data_path, filename), "wb") as fb:
                fb.write(new_data.astype("int16"))
    return None


def get_good_channels(data_path, amp_names):
    """
    Returns the size of the file in int16 bytes and the list of channels with that size.
    Inputs:
        data_path: the path to the data
        amp_names: the list of amp_names
    Returns:
        size: the size of the file in int16 bytes
        good_channels: the list of channels with that size
    """
    logging.debug("Function get_good_channels")

    filesize = {}
    for filename in amp_names:
        try:
            filesize[filename] = int(
                os.path.getsize(add_paths(data_path, filename)) / 2
            )
            # sizes.append(int(os.path.getsize(data_path+filename)/2))
        except:
            print(f"{filename} not exist.")
            filesize[filename] = 0
            # np.zeros(max(sizes), dtype=np.int16).tofile(data_path+filename)

    size = max(filesize.values())
    logging.debug(f"\t\t filesize: { max(filesize.values())}")
    if size == 0:
        return 0, None

    good_channels = [key for key in filesize.keys() if filesize[key] == size]
    logging.debug(f"\t\t number good channels: {len(good_channels)}")
    return size, good_channels


def single_file(data_path, amp_names, save_path, chunk_sz=512):
    """Create a NxM he simplest file format is the raw_binary one.
        Suppose you have N channels
        c0,c1,...,cN

        And if you assume that ci(t)
        is the value of channel ci

        at time t, then your datafile should be a raw file with values
        c0(0),c1(0),...,cN(0),c0(1),...,cN(1),...cN(T)

    This is simply the flatten version of your recordings matrix, with size N x T
    returns: bad channels

    Inputs:
        data_path: the path to the data
        amp_names: the list of amp_names
        save_path: the path to save the data
        chunk_sz: the size of the chunk in int16 bytes
    Returns:
        None, creates a new file .dat
    """
    print(f"{'-'*60}")
    print(f"Creating a single file from: {data_path}")

    out_name = add_paths(save_path, "".join(data_path.split(f"{os.sep}")[-3:]) + ".bin")
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    logging.debug(f"Started")
    try:
        info_file = add_paths(data_path, "info.rhd")
        logging.debug(f" info file : {info_file}")
        header = read_header(open(info_file, "rb"))
        sample_rate = np.asarray(header["sample_rate"])
        logging.debug(f" sample_rate : {sample_rate}")
        sample_rate.tofile(add_paths(save_path, f"sample_rate_{int(sample_rate)}.txt"))
    except:
        print("Could not determine sample rate")

    logging.debug(f"data path : {data_path}")
    logging.debug(f"amp_names : {amp_names}")
    size, good_channels = get_good_channels(data_path, amp_names)

    if not good_channels:
        logging.debug(f'sample_rate : {amp_names[0].split("-")[1]}')
        print(f'No data in channel {amp_names[0].split("-")[1]}')
        exit(-1)
    create_zeros(size, good_channels, data_path, amp_names)

    if size / chunk_sz > int(size / chunk_sz):
        n = int(size / chunk_sz) + 1
    else:
        n = int(size / chunk_sz)

    logging.debug(f"Start For loop")
    for i in tqdm(range(0, size, n)):
        chunk = np.array(
            [
                np.memmap(add_paths(data_path, filename), dtype="int16")[i : i + n]
                for filename in sorted(amp_names)
            ]
        ).T.flatten()

        if i == 0:
            with open(out_name, "wb") as fb:
                fb.write(chunk.astype("int16").tobytes())
        else:
            with open(out_name, "ab+") as fb:
                fb.write(chunk.astype("int16").tobytes())
    print("Done!")