From ba576396876a0ea5ab01a3b3cec67a81ebc0d6d5 Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Mon, 8 Nov 2021 15:23:44 -0800
Subject: [PATCH 1/7] add ecog data file readers to data.py

---
 aopy/data.py | 607 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 607 insertions(+)

diff --git a/aopy/data.py b/aopy/data.py
index a166e2ff..dfd48270 100644
--- a/aopy/data.py
+++ b/aopy/data.py
@@ -12,6 +12,15 @@
 import warnings
 import pickle
 
+import torch
+from torch.utils.data import Dataset, SubsetRandomSampler, RandomSampler, DataLoader
+import os.path as path # may need to build a switch here for PC/POSIX
+import re
+import json
+import pickle as pkl
+from torch.utils.data import dataset, IterableDataset
+import bisect
+
 def get_filenames_in_dir(base_dir, te):
     '''
     Gets the filenames for available systems in a given task entry. Requires that
@@ -1033,4 +1042,602 @@ def pkl_read(file_to_read, read_dir):
         this_dat = pickle.load(f)
     return this_dat
 
+# - - -- --- ----- -------- ------------- -------- ----- --- -- - - #
+# - - -- --- ----- -------- ------------- -------- ----- --- -- - - #
+
+class DataFile():
+    r''' DataFile() class - interface class for multichannel signal data stored in binary files. Allows for segment reading without full simultaneous RAM storage
+    inputs:
+        - data_file_path: string
+        - experiment_file_path=None: 
+        - mask_file_path=None:
+    
+    methods:
+        - read(): returns data segments defined by time start and stop points. Default behavior reads entire time span while masking channels as specified in data mask file.
+    '''
+
+    def __init__(self, data_file_path, exp_file_path=None, mask_file_path=None):
+
+        # parse file directory and components
+        data_dir = path.dirname(data_file_path)
+        data_basename = path.basename(data_file_path)
+        rec_id, device_id, rec_type, data_ext = data_basename.split('.')
+
+        # experiment data file: construct and load
+        if not exp_file_path:
+            exp_file_name = rec_id + 'experiment.json'
+            exp_file_path = path.join(data_dir,exp_file_name)
+
+        # mask file: construct and load
+        if not mask_file_path:
+            mask_file_name = rec_id + '.' + device_id + '.' + rec_type + '.mask.pkl'
+            mask_file_path = path.join(data_dir,mask_file_name)
+         
+        # set recording parameters
+        self.set_data_parameters(data_file_path,exp_file_path,mask_file_path)
+
+    # this is returned when the print() command is called.
+    def __repr__(self):
+        path_repr_str = f'Data file object: {self.data_file_path}'
+        sample_repr_str = f'\tsamples: {self.n_sample} ({self.n_sample/self.srate:0.2f}s, {self.data_mask.mean()*100:0.2f}% masked)'
+        ch_repr_str = f'\tchannels: {self.n_ch} ({self.ch_idx.mean()*100:0.2f}% masked)'
+        return path_repr_str + '\n' + sample_repr_str + '\n' + ch_repr_str + '\n'
+                
+
+    # read data segment. Default call (no arguments) returns the entire recording.
+    def read( self, t_start=0, t_len=-1, ch_idx=None, use_mask=True, mask_value=0., mask_pad_t=5 ):
+
+        # get offset sample/byte values
+        n_offset_samples = int(round(t_start * self.srate))
+        n_offset_items = n_offset_samples * self.n_ch
+        n_offset_bytes = n_offset_items * self.data_type().nbytes
+        if t_len == -1:
+            n_read_items = t_len
+            n_read_samples = int(self.n_sample)
+        else:
+            n_read_samples = int(t_len * self.srate)
+            n_read_items = n_read_samples * self.n_ch
+        
+        # read data
+        with open(self.data_file_path,'rb') as f:
+            data = np.fromfile(f,self.data_type,count=n_read_items,offset=n_offset_bytes)
+        data = np.reshape(data,(self.n_ch,n_read_samples),order=self.reshape_order)
+
+        # remove channels
+        if not ch_idx:
+            ch_idx = ~self.ch_idx
+        data = data[ch_idx,:] # mask values are True for bad spots
+
+        # mask data
+        sample_idx = np.arange(n_offset_samples,n_offset_samples+n_read_samples)
+        data[:,self.data_mask[sample_idx]] = mask_value
+
+        # consider: time array? May not want to incorporate until global time is added
+        return data
+
+    @staticmethod
+    def get_microdrive_parameters(exp_dict,microdrive_name):
+        microdrive_name_list = [md['name'] for md in exp_dict['hardware']['microdrive']]
+        microdrive_idx = [md_idx for md_idx, md in enumerate(microdrive_name_list) if microdrive_name == md][0]
+        microdrive_dict = exp_dict['hardware']['microdrive'][microdrive_idx]
+        electrode_label_list = [e['label'] for e in exp_dict['hardware']['microdrive'][0]['electrodes']]
+        n_ch = len(electrode_label_list)
+        return electrode_label_list, n_ch
+
+    @staticmethod
+    def get_read_parameters(exp_dict,rec_type):
+        clfp_pattern = 'clfp*'
+        if rec_type == 'raw':
+            srate = exp_dict['hardware']['acquisition']['samplingrate']
+            data_type = np.ushort
+            reshape_order = 'F'
+        elif rec_type == 'lfp':
+            srate = 1000
+            data_type = np.float32
+            reshape_order = 'F'
+        elif re.match(clfp_pattern,rec_type):
+            data_type = np.float32
+            if rec_type == 'clfp':
+                # there are a few different naming conventions, this is the default
+                srate = 1000
+                reshape_order = 'F'
+            else:
+                clfp_ds_pattern = 'clfp_ds(\d+)'
+                ds_match = re.search(clfp_ds_pattern,rec_type)
+                srate = int(ds_match.group(1))
+                reshape_order = 'C'
+        assert isinstance(srate,int), 'parsed srate value not an integer'
+        return srate, data_type, reshape_order
+
+    @staticmethod
+    def get_mask_file_path(data_path,rec_type,data_file_kern):
+        clfp_pattern = 'clfp*'
+        if rec_type == 'raw':
+            ecog_mask_file = None
+        elif rec_type == 'lfp':
+            ecog_mask_file = None
+        elif re.match(clfp_pattern,rec_type):
+            if rec_type == 'clfp':
+                ecog_mask_file = path.join(data_path,data_file_kern + ".mask.pkl")
+            else:
+                clfp_ds_pattern = 'clfp_ds(\d+)'
+                ds_match = re.search(clfp_ds_pattern,rec_type)
+                clfp_ds_file_kern = ".".join(data_file_kern.split(".")[:-1] + [ds_match.group()])
+                ecog_mask_file = path.join(data_path,clfp_ds_file_kern+".mask.pkl")
+        return ecog_mask_file
+                
+
+    # compute data parameter values and add as object attributes
+    def set_data_parameters( self, data_file_path, exp_file_path, mask_file_path):
+        # parse file
+        data_file = path.basename(data_file_path)
+        data_file_kern = path.splitext(data_file)[0]
+        rec_id, microdrive_name, rec_type = data_file_kern.split('.')
+        data_path = path.dirname(data_file_path)
+        
+        # read experiment file
+        exp_file = path.join(data_path,rec_id + ".experiment.json")
+        with open(exp_file,'r') as f:
+            exp_dict = json.load(f)
+        
+        # get microdrive parameters
+        electrode_label_list, n_ch = self.get_microdrive_parameters(exp_dict,microdrive_name)
+        
+        # get srate
+        srate, data_type, reshape_order = self.get_read_parameters(exp_dict, rec_type)
+        
+        # read mask
+        ecog_mask_file = self.get_mask_file_path(data_path,rec_type,data_file_kern)
+        with open(ecog_mask_file,"rb") as mask_f:
+            mask = pkl.load(mask_f)
+        # data_mask = grow_bool_array(mask["hf"] | mask["sat"], growth_size=int(srate*0.5))
+        data_mask = mask["hf"] | mask["sat"]
+        if 'ch' in mask.keys():
+            ch_idx = mask['ch']
+        else:
+            ch_idx = np.arange(n_ch)
+
+        # clean channel labels - formatting can change from recording to recording. Get Channel ID from full string.
+        ch_label_pattern = r'E\d+'
+        ch_label_cleaned = [re.findall(ch_label_pattern,ch_l)[0] for ch_l in electrode_label_list]
+        
+        # set parameters
+        self.data_file_path = data_file_path
+        self.exp_file_path = exp_file_path
+        self.mask_file_path = mask_file_path
+        self.rec_id = rec_id
+        self.microdrive_name = microdrive_name
+        self.rec_type = rec_type
+        self.srate = srate
+        self.data_type = data_type
+        self.reshape_order = reshape_order
+        self.data_mask = data_mask
+        self.n_ch = n_ch
+        self.ch_idx = ch_idx
+        self.ch_labels = ch_label_cleaned
+
+        # set sample length information
+        self.n_sample = len(self.data_mask)
+        self.t_total = self.n_sample/self.srate # (s)
+
+class DatafileDataset(Dataset):
+
+    r"""pytorch Dataset accessing Datafile interface.
+
+    Dataset object allowing (src, trg) sampling directly from structured binary data files.
+    Built to interface with aoLab datasets. Specifically constructed for the ECoG/LFP wireless platform data.
+
+    Arguments:
+        datafile (DataFile): DataFile object
+        src_t (float):\ttime length (s) of source sample
+        trg_t (float):\ttime length (s) of target sample
+        step_t (float):\ttime length (s) between src/trg pair sample starting points
+        transform (function):\tdata transformation method for adjusting sample output pairs.
+
+    """
+
+    def __init__( self, datafile, src_t, trg_t, step_t, in_mem=False, transform=None, device='cpu' ):
+        assert (isinstance(datafile, DataFile) or path.exists(datafile)), 'first argument must be DataFile object or valid path string'
+        if isinstance(datafile, str):
+            datafile = DataFile(datafile)
+        sample_t = src_t + trg_t
+        src_len = round(src_t*datafile.srate)
+        trg_len = round(trg_t*datafile.srate)
+        step_len = round(step_t*datafile.srate)
+        sample_len = round(sample_t*datafile.srate)
+        sample_start_idx = np.arange(0,datafile.n_sample-sample_len,step_len) # all candidate starting indices
+        sample_start_idx_in_mask = [np.any(datafile.data_mask[s_s_idx:s_s_idx+sample_len]) for s_s_idx in sample_start_idx] # sample window is masked
+        sample_start_idx = sample_start_idx[np.logical_not(sample_start_idx_in_mask)] # remove masked starting indices
+        sample_start_t = sample_start_idx/datafile.srate
+
+        # read whole data file if in_mem
+        if in_mem:
+            print(f'reading data from {datafile.data_file_path}...')
+            self.data = datafile.read()
+        else:
+            self.data = None
+
+        self.datafile = datafile
+        self.src_len = src_len
+        self.trg_len = trg_len
+        self.step_len = step_len
+        self.sample_len = sample_len
+        self.src_t = src_t
+        self.trg_t = trg_t
+        self.step_t = step_t
+        self.sample_t = sample_t
+        self.in_mem = in_mem
+        self.sample_start_idx = sample_start_idx
+        self.sample_start_t = sample_start_t
+        self.transform = transform
+        self.device = device
+
+    def read_sample( self, idx, ch_idx):
+        if self.in_mem:
+            sample_idx = np.arange(self.sample_len) + self.sample_start_idx[idx]
+            sample = self.data[:,sample_idx]
+            if ch_idx:
+                sample = sample[ch_idx,:]
+        else:
+            sample = self.datafile.read(t_start=self.sample_start_t[idx],t_len=self.sample_t,ch_idx=ch_idx)
+        return sample
+
+    def __len__( self ):
+        return len(self.sample_start_idx)
+
+    def __getitem__( self, idx, ch_idx=None ):
+        sample = self.read_sample(idx,ch_idx)
+        src = torch.tensor(sample[:,:self.src_len]).T
+        trg = torch.tensor(sample[:,self.src_len:]).T
+        if self.transform:
+            src,trg = self.transform((src,trg))
+        return src.to(self.device), trg.to(self.device)
+
+
+class DatafileConcatDataset(Dataset):
+    r"""Dataset as a concatenation of multiple datafile datasets.
+
+    This class is useful to assemble different existing datafile datasets and draw from the channel indices that they share.
+
+    Arguments:
+        datasets (sequence): List of datafile datasets to be concatenated
+    """
+
+    @staticmethod
+    def cumsum(sequence):
+        r, s = [], 0
+        for e in sequence:
+            l = len(e)
+            r.append(l + s)
+            s += l
+        return r
+
+    def __init__(self, datasets, transform=None):
+        super(DatafileConcatDataset, self).__init__()
+        assert len(datasets) > 0, 'datasets should not be an empty iterable'
+        self.datasets = list(datasets)
+        for d in self.datasets:
+            assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset"
+        self.cumulative_sizes = self.cumsum(self.datasets)
+        srate_set = list(set([ds.datafile.srate for ds in self.datasets]))
+        assert len(srate_set) == 1, 'all datasets must have the same sampling rate'
+        # get intersection of channel labels present in each dataset in self.datasets
+        file_labels = []
+        for d in self.datasets:
+            file_mask_idx = np.arange(d.datafile.n_ch)[~d.datafile.ch_idx] # idx of unmasked channels in this file
+            file_labels.append(np.array(d.datafile.ch_labels)[file_mask_idx])
+        self.ch_label = list(set(file_labels[0]).intersection(*file_labels))
+        # get index list of intersection channel locations in each datafile
+        ch_sample_idx_list = []
+        for d in self.datasets:
+            dataset_ch_sample_idx_list = []
+            for ch_i_l in self.ch_label:
+                dataset_ch_sample_idx_list.append(list(np.array(d.datafile.ch_labels)[~np.array(d.datafile.ch_idx)]).index(ch_i_l))
+            ch_sample_idx_list.append(dataset_ch_sample_idx_list)
+        self.ch_idx = ch_sample_idx_list
+        self.n_ch = len(self.ch_label)
+        self.srate = srate_set[0]
+        self.transform = transform
+
+    def __len__(self):
+        return self.cumulative_sizes[-1]
+
+    def __getitem__(self, idx):
+        if idx < 0:
+            if -idx > len(self):
+                raise ValueError("absolute value of index should not exceed dataset length")
+            idx = len(self) + idx
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        if dataset_idx == 0:
+            sample_idx = idx
+        else:
+            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+        src, trg = self.datasets[dataset_idx].__getitem__(sample_idx)
+        src = src[:,self.ch_idx[dataset_idx]]
+        trg = trg[:,self.ch_idx[dataset_idx]]
+        if self.transform:
+            src, trg = self.transform((src,trg))
+        return src, trg
+
+    def get_data_loaders( self, partition=(0.8,0.2,0.0), batch_size=1, num_workers=0, rand_part=False, rand_seed=42 ):
+        r'''
+            Return dataloader objects for accessing training, validation and testing 
+            partitions of the DatafileConcatDataset. Dataloaders can sample sequentially or randomly.
+
+            arguments:
+                - partition (default (0.8,0.2,0.0)): tuple of partition fractional sizes (train_frac, valid_frac, test_frac).
+                    Values will be normalized to sum to 1.
+                - batch_size (default 1): int value defining the size of each batch draw
+                - rand_part (default False): bool determining sequential or random partitioning
+                - rand_seed (default 42): int setting the rng. Keeps partitions consistent
+        '''
+        # get partition sizes
+        frac_sum = np.sum(partition)
+        train_frac = partition[0]/frac_sum
+        valid_frac = partition[1]/frac_sum
+        test_frac = partition[2]/frac_sum
+        n_train_samp = round(train_frac * len(self))
+        n_valid_samp = round(valid_frac * len(self))
+        n_test_samp = round(test_frac * len(self))
+        # create partition index arrays
+        if rand_part:
+            if not isinstance(rand_seed, int):
+                try: rand_seed_new = int(rand_seed)
+                except:
+                    raise TypeError(f'Could not cast rand_seed value {rand_seed} to int.')
+                raise Warning(f'ValueWarning: rand_seed must be of type int. Casting from {type(rand_seed)} {rand_seed} to int {rand_seed_new}. This might cause issues.')
+            sample_idx = np.random.RandomState(seed=rand_seed).permutation(len(self))
+        else:
+            sample_idx = np.arange(len(self))
+        train_sample_idx = sample_idx[:n_train_samp]
+        valid_sample_idx = sample_idx[n_train_samp:(n_train_samp+n_valid_samp)]
+        test_sample_idx = sample_idx[(n_train_samp+n_valid_samp):]
+        # create samplers
+        train_sampler = SubsetRandomSampler(train_sample_idx)
+        valid_sampler = SubsetRandomSampler(valid_sample_idx)
+        test_sampler = SubsetRandomSampler(test_sample_idx)
+        # create dataloaders
+        train_loader = DataLoader(self,batch_size=batch_size,sampler=train_sampler,num_workers=num_workers)
+        valid_loader = DataLoader(self,batch_size=batch_size,sampler=valid_sampler,num_workers=num_workers)
+        test_loader = DataLoader(self,batch_size=batch_size,sampler=test_sampler,num_workers=num_workers)
+
+        return train_loader, valid_loader, test_loader
+
+    @property
+    def cummulative_sizes( self ):
+        warnings.warn("cummulative_sizes attribute is renamed to "
+                      "cumulative_sizes", DeprecationWarning, stacklevel=2)
+        return self.cumulative_sizes
+
+def data_transform_normalize( src, trg, scale_factor=1. ):
+    r'''Data transform. Normalizes src, trg pairs through z-scoring.
+    '''
+
+    sample = np.concatenate((src,trg),axis=-1)
+    center = np.mean(sample,axis=-1)
+    std = np.std(sample,axis=-1)
+    src = scale_factor * ((src.T - center)/std).T # is there a better way to align dimensions? einsum?
+    trg = scale_factor * ((trg.T - center)/std).T
+    return (src, trg)
+
+
+def parse_file_info(file_path):
+    file_name = os.path.basename(file_path)
+    data_file_noext = os.path.splitext(file_name)[0]
+    data_file_parts = data_file_noext.split('.')
+    if len(data_file_parts) == 3:
+        rec_id, microdrive_name, rec_type = data_file_parts
+    else:
+        rec_id, microdrive_name, _, rec_type = data_file_parts
+    data_dir = os.path.dirname(file_path)
+    exp_file_name = os.path.join(data_dir,rec_id + ".experiment.json")
+    mask_file_name = os.path.join(data_dir,data_file_noext + ".mask.pkl")
+    return exp_file_name, mask_file_name, microdrive_name, rec_type
+
+def load_experiment_data(exp_file_name):
+    assert os.path.exists(exp_file_name), f'inferred experiment file not found at {exp_file_name}'
+    with open(exp_file_name,'r') as f:
+        experiment = json.load(f)
+    electrode_df = DataFrame(experiment['hardware']['microdrive'][0]['electrodes'])
+    electrode_df = DataFrame.join(electrode_df,DataFrame(list(electrode_df.position)))
+    del electrode_df['position']
+    return experiment, electrode_df
+
+def load_mask_data(mask_file_name):
+    assert os.path.exists(mask_file_name), f'inferred mask file not found at {mask_file_name}'
+    with open(mask_file_name,'r') as f:
+        mask = pkl.load(f)
+
+def read_lfp(file_path,t_range=(0,-1)):
+
+    # get local experiment, mask files
+    exp_file_name, mask_file_name, microdrive_name, rec_type = parse_file_info(file_path)
+
+    # load experiment data
+    experiment, electrode_df = load_experiment_data(exp_file_name)
+
+    # load mask data
+    mask = load_mask_data(mask_file_name)
+
+    # get parameters: srate
+    dsmatch = re.search('clfp_ds(\d+)',rec_type)
+    if rec_type == 'raw':
+        srate = experiment['hardware']['acquisition']['samplingrate']
+        data_type = np.ushort
+        reshape_order = 'F'
+    elif rec_type == 'lfp':
+        srate = 1000
+        data_type = np.float32
+        reshape_order = 'F'
+    elif rec_type == 'clfp':
+        srate = 1000
+        data_type = np.float32
+        reshape_order = 'F'
+    elif dsmatch:
+        # downsampled data - get srate from name
+        srate = int(dsmatch.group(1))
+        data_type = np.float32
+        reshape_order = 'C' # files created with np.tofile which forces C ordering.
+
+    # get microdrive parameters
+    microdrive_name_list = [md['name'] for md in experiment['hardware']['microdrive']]
+    microdrive_idx = [md_idx for md_idx, md in enumerate(microdrive_name_list) if microdrive_name == md][0]
+    microdrive_dict = experiment['hardware']['microdrive'][microdrive_idx]
+    num_ch = len(microdrive_dict['electrodes'])
+
+    # get file size information
+    data_type_size = data_type().nbytes
+    file_size = os.path.getsize(file_path)
+    n_offset_samples = np.round(t_range[0]*srate)
+    n_offset_bytes = n_offset_samples*data_type_size
+    n_all = int(np.floor(file_size/num_ch/data_type_size))
+    n_stop = n_all if t_range[1] == -1 else np.min((np.round(t_range[1]*srate),n_all))
+    n_read = n_stop-n_offset_samples
+
+    # read signal data
+    data = read_from_file(
+        file_path,
+        data_type,
+        num_ch,
+        n_read,
+        n_offset_bytes,
+        reshape_order=reshape_order
+    )
+
+    # create xarray from data and channel information
+    da = xr.DataArray(
+        data.T,
+        dime = ('sample','ch'),
+        coords = {
+            'ch': electrode_df.label,
+            'x_pos': ('ch', electrode_df.x),
+            'y_pos': ('ch', electrode_df.y),
+            'row': ('ch', electrode_df.row),
+            'col': ('ch', electrode_df.col),
+        },
+        attrs = {'srate': srate}
+    )
+
+    return da, mask
+
+# wrapper to read and handle clfp ECOG data
+def load_ecog_clfp_data(data_file_name,t_range=(0,-1),exp_file_name=None,mask_file_name=None,compute_mask=True):
+
+    # get file path, set ancillary data file names
+    exp_file_name, mask_file_name, microdrive_name, rec_type = parse_file_info(data_file_name)
+
+    # check for experiment file, load if valid, exit if not.
+    if os.path.exists(exp_file_name):
+        with open(exp_file_name,'r') as f:
+            experiment = json.load(f)
+    else:
+        raise NameError(f'Experiment file {exp_file_name} either invalid or not found. Aborting Process.')
+
+    # get srate
+    dsmatch = re.search('clfp_ds(\d+)',rec_type)
+    if rec_type == 'raw':
+        srate = experiment['hardware']['acquisition']['samplingrate']
+        data_type = np.ushort
+        reshape_order = 'F'
+    elif rec_type == 'lfp':
+        srate = 1000
+        data_type = np.float32
+        reshape_order = 'F'
+    elif rec_type == 'clfp':
+        srate = 1000
+        data_type = np.float32
+        reshape_order = 'F'
+    elif dsmatch:
+        # downsampled data - get srate from name
+        srate = int(dsmatch.group(1))
+        data_type = np.float32
+        compute_mask = False
+        reshape_order = 'C' # files created with np.tofile which forces C ordering. Sorry!
+    else:
+        raise NameError(f'File type {rec_type}.dat not recognized. Aborting read process.')
+
+    # get microdrive parameters
+    microdrive_name_list = [md['name'] for md in experiment['hardware']['microdrive']]
+    microdrive_idx = [md_idx for md_idx, md in enumerate(microdrive_name_list) if microdrive_name == md][0]
+    microdrive_dict = experiment['hardware']['microdrive'][microdrive_idx]
+    num_ch = len(microdrive_dict['electrodes'])
+
+    exp = {"srate":srate,"num_ch":num_ch}
+
+    data_type_size = data_type().nbytes
+    file_size = os.path.getsize(data_file_name)
+    n_offset_samples = np.round(t_range[0]*srate)
+    n_offset = n_offset_samples*data_type_size
+    n_all = int(np.floor(file_size/num_ch/data_type_size))
+    if t_range[1] == -1:
+        n_stop = n_all
+    else:
+        n_stop = np.min((np.round(t_range[1]*srate),n_all))
+    n_read = n_stop-n_offset_samples
+
+    # load data
+    print("Loading data file:")
+    # n_offset value is the number of bytes to skip
+    # n_read value is the number of items to read (by data type)
+    data = read_from_file(data_file_name,data_type,num_ch,n_read,n_offset,
+                          reshape_order=reshape_order)
+    if rec_type == 'raw': # correct uint16 encoding errors
+        data = np.array(data,dtype=np.float32)
+        for ch_idx in range(num_ch):
+            is_neg = data[ch_idx,:] > 2**15
+            data[ch_idx,is_neg] = data[ch_idx,is_neg] - (2**16 - 1)
+
+    # check for mask file, load if valid, compute if not
+    if os.path.exists(mask_file_name):
+        with open(mask_file_name,"rb") as mask_f:
+            mask = pkl.load(mask_f)
+    elif compute_mask:
+        print("No mask data file found for {0}".format(data_file))
+        print("Computing data masks:")
+        hf_mask,_ = datafilter.high_freq_data_detection(data,srate)
+        _,sat_mask_all = datafilter.saturated_data_detection(data,srate)
+        sat_mask = np.any(sat_mask_all,axis=0)
+        mask = {"hf":hf_mask,"sat":sat_mask}
+        # save mask data to current directory
+        print("Saving mask data for {0} to {1}".format(data_file,mask_file_name))
+        with open(mask_file_name,"wb") as mask_f:
+            pkl.dump(mask,mask_f)
+    else:
+        mask = []
+
+    return data, exp, mask
+
+# read T seconds of data from the start of the recording:
+def read_from_start(data_file_path,data_type,n_ch,n_read):
+    data_file = open(data_file_path,"rb")
+    data = np.fromfile(data_file,dtype=data_type,count=n_read*n_ch)
+    data = np.reshape(data,(n_ch,n_read),order='F')
+    data_file.close()
+
+    return data
+
+# read some time from a given offset
+def read_from_file(data_file_path,data_type,n_ch,n_read,n_offset,reshape_order='F'):
+    data_file = open(data_file_path,"rb")
+    if np.version.version >= "1.17": # "offset" field not added until later installations
+        data = np.fromfile(data_file,dtype=data_type,count=n_read*n_ch,
+                           offset=n_offset*n_ch)
+    else:
+        warnings.warn("'offset' feature not available in numpy <= 1.13 - reading from the top",FutureWarning)
+        data = np.fromfile(data_file,dtype=data_type,count=n_read*n_ch)
+    data = np.reshape(data,(n_ch,n_read),order=reshape_order)
+    data_file.close()
+
+    return data
+
+# read variables from the "experiment.mat" files
+def get_exp_var(exp_data,*args):
+    out = exp_data.copy()
+    for k, var_name in enumerate(args):
+        if k > 1:
+            out = out[None][0][None][0][var_name]
+
+        else:
+            out = out[var_name]
 
+    return out

From 988572db247f65c9e7e6671d860a468eef4379dc Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Mon, 8 Nov 2021 15:27:37 -0800
Subject: [PATCH 2/7] Enh: clean dependencies in data.py

---
 aopy/data.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/aopy/data.py b/aopy/data.py
index dfd48270..8363dc5d 100644
--- a/aopy/data.py
+++ b/aopy/data.py
@@ -6,12 +6,10 @@
 import h5py
 import tables
 import csv
-import pandas as pd
 import os
 import glob
 import warnings
 import pickle
-
 import torch
 from torch.utils.data import Dataset, SubsetRandomSampler, RandomSampler, DataLoader
 import os.path as path # may need to build a switch here for PC/POSIX
@@ -20,6 +18,10 @@
 import pickle as pkl
 from torch.utils.data import dataset, IterableDataset
 import bisect
+import numpy as np
+from pandas import read_csv, read_excel, DataFrame
+import xarray as xr
+import warnings
 
 def get_filenames_in_dir(base_dir, te):
     '''
@@ -158,8 +160,8 @@ def load_optitrack_data(data_dir, filename):
     filepath = os.path.join(data_dir, filename)
     # Load .csv file as a pandas data frame, convert to a numpy array, and remove
     # the 'Frame' and 'Time (Seconds)' columns.
-    mocap_data_rot = pd.read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,mocap_data_rot_column_idx]
-    mocap_data_pos = pd.read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,mocap_data_pos_column_idx]
+    mocap_data_rot = read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,mocap_data_rot_column_idx]
+    mocap_data_pos = read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,mocap_data_pos_column_idx]
 
     return mocap_data_pos, mocap_data_rot
 
@@ -182,7 +184,7 @@ def load_optitrack_time(data_dir, filename):
     filepath = os.path.join(data_dir, filename)
     # Load .csv file as a pandas data frame, convert to a numpy array, and only
     # return the 'Time (Seconds)' column
-    timestamps = pd.read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,timestamp_column_idx]
+    timestamps = read_csv(filepath, header=column_names_idx_csvrow).to_numpy()[:,timestamp_column_idx]
     return timestamps
 
 
@@ -706,7 +708,7 @@ def lookup_excel_value(data_dir, excel_file, from_column, to_column, lookup_valu
     if fullfile in _cached_dataframes:
         dataframe = _cached_dataframes[fullfile]
     else:
-        dataframe = pd.read_excel(fullfile)
+        dataframe = read_excel(fullfile)
         _cached_dataframes[fullfile] = dataframe
     
     row = dataframe.loc[dataframe[from_column] == lookup_value]
@@ -762,7 +764,7 @@ def load_electrode_pos(data_dir, pos_file):
             | **y_pos (nch):** y position of each electrode
     '''
     fullfile = os.path.join(data_dir, pos_file)
-    electrode_pos = pd.read_excel(fullfile)
+    electrode_pos = read_excel(fullfile)
     x_pos = electrode_pos['topdown_x'].to_numpy()
     y_pos = electrode_pos['topdown_y'].to_numpy()
     return x_pos, y_pos

From ace3048d51f9c621c1f524f6f46162f4f87dce32 Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Mon, 8 Nov 2021 15:31:38 -0800
Subject: [PATCH 3/7] Fix: edit excape character formatting in data.py

---
 aopy/data.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/aopy/data.py b/aopy/data.py
index 8363dc5d..9761e539 100644
--- a/aopy/data.py
+++ b/aopy/data.py
@@ -1162,7 +1162,7 @@ def get_mask_file_path(data_path,rec_type,data_file_kern):
             if rec_type == 'clfp':
                 ecog_mask_file = path.join(data_path,data_file_kern + ".mask.pkl")
             else:
-                clfp_ds_pattern = 'clfp_ds(\d+)'
+                clfp_ds_pattern = r'clfp_ds(\d+)'
                 ds_match = re.search(clfp_ds_pattern,rec_type)
                 clfp_ds_file_kern = ".".join(data_file_kern.split(".")[:-1] + [ds_match.group()])
                 ecog_mask_file = path.join(data_path,clfp_ds_file_kern+".mask.pkl")
@@ -1462,7 +1462,7 @@ def read_lfp(file_path,t_range=(0,-1)):
     mask = load_mask_data(mask_file_name)
 
     # get parameters: srate
-    dsmatch = re.search('clfp_ds(\d+)',rec_type)
+    dsmatch = re.search(r'clfp_ds(\d+)',rec_type)
     if rec_type == 'raw':
         srate = experiment['hardware']['acquisition']['samplingrate']
         data_type = np.ushort
@@ -1536,7 +1536,7 @@ def load_ecog_clfp_data(data_file_name,t_range=(0,-1),exp_file_name=None,mask_fi
         raise NameError(f'Experiment file {exp_file_name} either invalid or not found. Aborting Process.')
 
     # get srate
-    dsmatch = re.search('clfp_ds(\d+)',rec_type)
+    dsmatch = re.search(r'clfp_ds(\d+)',rec_type)
     if rec_type == 'raw':
         srate = experiment['hardware']['acquisition']['samplingrate']
         data_type = np.ushort

From 0b87dabcfc32c1a5f7be2c97f8c9dbbbc01a3af7 Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Mon, 8 Nov 2021 15:34:28 -0800
Subject: [PATCH 4/7] Fix: add xarray package dependency

---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index c1d11393..eeb99123 100644
--- a/setup.py
+++ b/setup.py
@@ -8,6 +8,7 @@
 
 install_requires = [
     'numpy',
+    'xarray',
     'pandas',
     'psutil',
     'h5py',

From 31ace01f3c35a031cc2680a6f60dd9538e94013d Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Mon, 8 Nov 2021 15:42:46 -0800
Subject: [PATCH 5/7] Remove pytorch code additions to data.py; the Dataset()
 class was overwriting the whitematter Dataset() class. I'll rework this.

---
 aopy/data.py | 384 +--------------------------------------------------
 1 file changed, 2 insertions(+), 382 deletions(-)

diff --git a/aopy/data.py b/aopy/data.py
index 9761e539..9c519c5e 100644
--- a/aopy/data.py
+++ b/aopy/data.py
@@ -9,15 +9,10 @@
 import os
 import glob
 import warnings
-import pickle
-import torch
-from torch.utils.data import Dataset, SubsetRandomSampler, RandomSampler, DataLoader
 import os.path as path # may need to build a switch here for PC/POSIX
 import re
 import json
 import pickle as pkl
-from torch.utils.data import dataset, IterableDataset
-import bisect
 import numpy as np
 from pandas import read_csv, read_excel, DataFrame
 import xarray as xr
@@ -1024,7 +1019,7 @@ def pkl_write(file_to_write, values_to_dump, write_dir):
     '''
     file = os.path.join(write_dir, file_to_write)
     with open(file, 'wb') as pickle_file:
-        pickle.dump(values_to_dump, pickle_file)
+        pkl.dump(values_to_dump, pickle_file)
 
 
 def pkl_read(file_to_read, read_dir):
@@ -1041,387 +1036,12 @@ def pkl_read(file_to_read, read_dir):
     '''
     file = os.path.join(read_dir, file_to_read)
     with open(file, "rb") as f:
-        this_dat = pickle.load(f)
+        this_dat = pkl.load(f)
     return this_dat
 
 # - - -- --- ----- -------- ------------- -------- ----- --- -- - - #
 # - - -- --- ----- -------- ------------- -------- ----- --- -- - - #
 
-class DataFile():
-    r''' DataFile() class - interface class for multichannel signal data stored in binary files. Allows for segment reading without full simultaneous RAM storage
-    inputs:
-        - data_file_path: string
-        - experiment_file_path=None: 
-        - mask_file_path=None:
-    
-    methods:
-        - read(): returns data segments defined by time start and stop points. Default behavior reads entire time span while masking channels as specified in data mask file.
-    '''
-
-    def __init__(self, data_file_path, exp_file_path=None, mask_file_path=None):
-
-        # parse file directory and components
-        data_dir = path.dirname(data_file_path)
-        data_basename = path.basename(data_file_path)
-        rec_id, device_id, rec_type, data_ext = data_basename.split('.')
-
-        # experiment data file: construct and load
-        if not exp_file_path:
-            exp_file_name = rec_id + 'experiment.json'
-            exp_file_path = path.join(data_dir,exp_file_name)
-
-        # mask file: construct and load
-        if not mask_file_path:
-            mask_file_name = rec_id + '.' + device_id + '.' + rec_type + '.mask.pkl'
-            mask_file_path = path.join(data_dir,mask_file_name)
-         
-        # set recording parameters
-        self.set_data_parameters(data_file_path,exp_file_path,mask_file_path)
-
-    # this is returned when the print() command is called.
-    def __repr__(self):
-        path_repr_str = f'Data file object: {self.data_file_path}'
-        sample_repr_str = f'\tsamples: {self.n_sample} ({self.n_sample/self.srate:0.2f}s, {self.data_mask.mean()*100:0.2f}% masked)'
-        ch_repr_str = f'\tchannels: {self.n_ch} ({self.ch_idx.mean()*100:0.2f}% masked)'
-        return path_repr_str + '\n' + sample_repr_str + '\n' + ch_repr_str + '\n'
-                
-
-    # read data segment. Default call (no arguments) returns the entire recording.
-    def read( self, t_start=0, t_len=-1, ch_idx=None, use_mask=True, mask_value=0., mask_pad_t=5 ):
-
-        # get offset sample/byte values
-        n_offset_samples = int(round(t_start * self.srate))
-        n_offset_items = n_offset_samples * self.n_ch
-        n_offset_bytes = n_offset_items * self.data_type().nbytes
-        if t_len == -1:
-            n_read_items = t_len
-            n_read_samples = int(self.n_sample)
-        else:
-            n_read_samples = int(t_len * self.srate)
-            n_read_items = n_read_samples * self.n_ch
-        
-        # read data
-        with open(self.data_file_path,'rb') as f:
-            data = np.fromfile(f,self.data_type,count=n_read_items,offset=n_offset_bytes)
-        data = np.reshape(data,(self.n_ch,n_read_samples),order=self.reshape_order)
-
-        # remove channels
-        if not ch_idx:
-            ch_idx = ~self.ch_idx
-        data = data[ch_idx,:] # mask values are True for bad spots
-
-        # mask data
-        sample_idx = np.arange(n_offset_samples,n_offset_samples+n_read_samples)
-        data[:,self.data_mask[sample_idx]] = mask_value
-
-        # consider: time array? May not want to incorporate until global time is added
-        return data
-
-    @staticmethod
-    def get_microdrive_parameters(exp_dict,microdrive_name):
-        microdrive_name_list = [md['name'] for md in exp_dict['hardware']['microdrive']]
-        microdrive_idx = [md_idx for md_idx, md in enumerate(microdrive_name_list) if microdrive_name == md][0]
-        microdrive_dict = exp_dict['hardware']['microdrive'][microdrive_idx]
-        electrode_label_list = [e['label'] for e in exp_dict['hardware']['microdrive'][0]['electrodes']]
-        n_ch = len(electrode_label_list)
-        return electrode_label_list, n_ch
-
-    @staticmethod
-    def get_read_parameters(exp_dict,rec_type):
-        clfp_pattern = 'clfp*'
-        if rec_type == 'raw':
-            srate = exp_dict['hardware']['acquisition']['samplingrate']
-            data_type = np.ushort
-            reshape_order = 'F'
-        elif rec_type == 'lfp':
-            srate = 1000
-            data_type = np.float32
-            reshape_order = 'F'
-        elif re.match(clfp_pattern,rec_type):
-            data_type = np.float32
-            if rec_type == 'clfp':
-                # there are a few different naming conventions, this is the default
-                srate = 1000
-                reshape_order = 'F'
-            else:
-                clfp_ds_pattern = 'clfp_ds(\d+)'
-                ds_match = re.search(clfp_ds_pattern,rec_type)
-                srate = int(ds_match.group(1))
-                reshape_order = 'C'
-        assert isinstance(srate,int), 'parsed srate value not an integer'
-        return srate, data_type, reshape_order
-
-    @staticmethod
-    def get_mask_file_path(data_path,rec_type,data_file_kern):
-        clfp_pattern = 'clfp*'
-        if rec_type == 'raw':
-            ecog_mask_file = None
-        elif rec_type == 'lfp':
-            ecog_mask_file = None
-        elif re.match(clfp_pattern,rec_type):
-            if rec_type == 'clfp':
-                ecog_mask_file = path.join(data_path,data_file_kern + ".mask.pkl")
-            else:
-                clfp_ds_pattern = r'clfp_ds(\d+)'
-                ds_match = re.search(clfp_ds_pattern,rec_type)
-                clfp_ds_file_kern = ".".join(data_file_kern.split(".")[:-1] + [ds_match.group()])
-                ecog_mask_file = path.join(data_path,clfp_ds_file_kern+".mask.pkl")
-        return ecog_mask_file
-                
-
-    # compute data parameter values and add as object attributes
-    def set_data_parameters( self, data_file_path, exp_file_path, mask_file_path):
-        # parse file
-        data_file = path.basename(data_file_path)
-        data_file_kern = path.splitext(data_file)[0]
-        rec_id, microdrive_name, rec_type = data_file_kern.split('.')
-        data_path = path.dirname(data_file_path)
-        
-        # read experiment file
-        exp_file = path.join(data_path,rec_id + ".experiment.json")
-        with open(exp_file,'r') as f:
-            exp_dict = json.load(f)
-        
-        # get microdrive parameters
-        electrode_label_list, n_ch = self.get_microdrive_parameters(exp_dict,microdrive_name)
-        
-        # get srate
-        srate, data_type, reshape_order = self.get_read_parameters(exp_dict, rec_type)
-        
-        # read mask
-        ecog_mask_file = self.get_mask_file_path(data_path,rec_type,data_file_kern)
-        with open(ecog_mask_file,"rb") as mask_f:
-            mask = pkl.load(mask_f)
-        # data_mask = grow_bool_array(mask["hf"] | mask["sat"], growth_size=int(srate*0.5))
-        data_mask = mask["hf"] | mask["sat"]
-        if 'ch' in mask.keys():
-            ch_idx = mask['ch']
-        else:
-            ch_idx = np.arange(n_ch)
-
-        # clean channel labels - formatting can change from recording to recording. Get Channel ID from full string.
-        ch_label_pattern = r'E\d+'
-        ch_label_cleaned = [re.findall(ch_label_pattern,ch_l)[0] for ch_l in electrode_label_list]
-        
-        # set parameters
-        self.data_file_path = data_file_path
-        self.exp_file_path = exp_file_path
-        self.mask_file_path = mask_file_path
-        self.rec_id = rec_id
-        self.microdrive_name = microdrive_name
-        self.rec_type = rec_type
-        self.srate = srate
-        self.data_type = data_type
-        self.reshape_order = reshape_order
-        self.data_mask = data_mask
-        self.n_ch = n_ch
-        self.ch_idx = ch_idx
-        self.ch_labels = ch_label_cleaned
-
-        # set sample length information
-        self.n_sample = len(self.data_mask)
-        self.t_total = self.n_sample/self.srate # (s)
-
-class DatafileDataset(Dataset):
-
-    r"""pytorch Dataset accessing Datafile interface.
-
-    Dataset object allowing (src, trg) sampling directly from structured binary data files.
-    Built to interface with aoLab datasets. Specifically constructed for the ECoG/LFP wireless platform data.
-
-    Arguments:
-        datafile (DataFile): DataFile object
-        src_t (float):\ttime length (s) of source sample
-        trg_t (float):\ttime length (s) of target sample
-        step_t (float):\ttime length (s) between src/trg pair sample starting points
-        transform (function):\tdata transformation method for adjusting sample output pairs.
-
-    """
-
-    def __init__( self, datafile, src_t, trg_t, step_t, in_mem=False, transform=None, device='cpu' ):
-        assert (isinstance(datafile, DataFile) or path.exists(datafile)), 'first argument must be DataFile object or valid path string'
-        if isinstance(datafile, str):
-            datafile = DataFile(datafile)
-        sample_t = src_t + trg_t
-        src_len = round(src_t*datafile.srate)
-        trg_len = round(trg_t*datafile.srate)
-        step_len = round(step_t*datafile.srate)
-        sample_len = round(sample_t*datafile.srate)
-        sample_start_idx = np.arange(0,datafile.n_sample-sample_len,step_len) # all candidate starting indices
-        sample_start_idx_in_mask = [np.any(datafile.data_mask[s_s_idx:s_s_idx+sample_len]) for s_s_idx in sample_start_idx] # sample window is masked
-        sample_start_idx = sample_start_idx[np.logical_not(sample_start_idx_in_mask)] # remove masked starting indices
-        sample_start_t = sample_start_idx/datafile.srate
-
-        # read whole data file if in_mem
-        if in_mem:
-            print(f'reading data from {datafile.data_file_path}...')
-            self.data = datafile.read()
-        else:
-            self.data = None
-
-        self.datafile = datafile
-        self.src_len = src_len
-        self.trg_len = trg_len
-        self.step_len = step_len
-        self.sample_len = sample_len
-        self.src_t = src_t
-        self.trg_t = trg_t
-        self.step_t = step_t
-        self.sample_t = sample_t
-        self.in_mem = in_mem
-        self.sample_start_idx = sample_start_idx
-        self.sample_start_t = sample_start_t
-        self.transform = transform
-        self.device = device
-
-    def read_sample( self, idx, ch_idx):
-        if self.in_mem:
-            sample_idx = np.arange(self.sample_len) + self.sample_start_idx[idx]
-            sample = self.data[:,sample_idx]
-            if ch_idx:
-                sample = sample[ch_idx,:]
-        else:
-            sample = self.datafile.read(t_start=self.sample_start_t[idx],t_len=self.sample_t,ch_idx=ch_idx)
-        return sample
-
-    def __len__( self ):
-        return len(self.sample_start_idx)
-
-    def __getitem__( self, idx, ch_idx=None ):
-        sample = self.read_sample(idx,ch_idx)
-        src = torch.tensor(sample[:,:self.src_len]).T
-        trg = torch.tensor(sample[:,self.src_len:]).T
-        if self.transform:
-            src,trg = self.transform((src,trg))
-        return src.to(self.device), trg.to(self.device)
-
-
-class DatafileConcatDataset(Dataset):
-    r"""Dataset as a concatenation of multiple datafile datasets.
-
-    This class is useful to assemble different existing datafile datasets and draw from the channel indices that they share.
-
-    Arguments:
-        datasets (sequence): List of datafile datasets to be concatenated
-    """
-
-    @staticmethod
-    def cumsum(sequence):
-        r, s = [], 0
-        for e in sequence:
-            l = len(e)
-            r.append(l + s)
-            s += l
-        return r
-
-    def __init__(self, datasets, transform=None):
-        super(DatafileConcatDataset, self).__init__()
-        assert len(datasets) > 0, 'datasets should not be an empty iterable'
-        self.datasets = list(datasets)
-        for d in self.datasets:
-            assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset"
-        self.cumulative_sizes = self.cumsum(self.datasets)
-        srate_set = list(set([ds.datafile.srate for ds in self.datasets]))
-        assert len(srate_set) == 1, 'all datasets must have the same sampling rate'
-        # get intersection of channel labels present in each dataset in self.datasets
-        file_labels = []
-        for d in self.datasets:
-            file_mask_idx = np.arange(d.datafile.n_ch)[~d.datafile.ch_idx] # idx of unmasked channels in this file
-            file_labels.append(np.array(d.datafile.ch_labels)[file_mask_idx])
-        self.ch_label = list(set(file_labels[0]).intersection(*file_labels))
-        # get index list of intersection channel locations in each datafile
-        ch_sample_idx_list = []
-        for d in self.datasets:
-            dataset_ch_sample_idx_list = []
-            for ch_i_l in self.ch_label:
-                dataset_ch_sample_idx_list.append(list(np.array(d.datafile.ch_labels)[~np.array(d.datafile.ch_idx)]).index(ch_i_l))
-            ch_sample_idx_list.append(dataset_ch_sample_idx_list)
-        self.ch_idx = ch_sample_idx_list
-        self.n_ch = len(self.ch_label)
-        self.srate = srate_set[0]
-        self.transform = transform
-
-    def __len__(self):
-        return self.cumulative_sizes[-1]
-
-    def __getitem__(self, idx):
-        if idx < 0:
-            if -idx > len(self):
-                raise ValueError("absolute value of index should not exceed dataset length")
-            idx = len(self) + idx
-        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
-        if dataset_idx == 0:
-            sample_idx = idx
-        else:
-            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
-        src, trg = self.datasets[dataset_idx].__getitem__(sample_idx)
-        src = src[:,self.ch_idx[dataset_idx]]
-        trg = trg[:,self.ch_idx[dataset_idx]]
-        if self.transform:
-            src, trg = self.transform((src,trg))
-        return src, trg
-
-    def get_data_loaders( self, partition=(0.8,0.2,0.0), batch_size=1, num_workers=0, rand_part=False, rand_seed=42 ):
-        r'''
-            Return dataloader objects for accessing training, validation and testing 
-            partitions of the DatafileConcatDataset. Dataloaders can sample sequentially or randomly.
-
-            arguments:
-                - partition (default (0.8,0.2,0.0)): tuple of partition fractional sizes (train_frac, valid_frac, test_frac).
-                    Values will be normalized to sum to 1.
-                - batch_size (default 1): int value defining the size of each batch draw
-                - rand_part (default False): bool determining sequential or random partitioning
-                - rand_seed (default 42): int setting the rng. Keeps partitions consistent
-        '''
-        # get partition sizes
-        frac_sum = np.sum(partition)
-        train_frac = partition[0]/frac_sum
-        valid_frac = partition[1]/frac_sum
-        test_frac = partition[2]/frac_sum
-        n_train_samp = round(train_frac * len(self))
-        n_valid_samp = round(valid_frac * len(self))
-        n_test_samp = round(test_frac * len(self))
-        # create partition index arrays
-        if rand_part:
-            if not isinstance(rand_seed, int):
-                try: rand_seed_new = int(rand_seed)
-                except:
-                    raise TypeError(f'Could not cast rand_seed value {rand_seed} to int.')
-                raise Warning(f'ValueWarning: rand_seed must be of type int. Casting from {type(rand_seed)} {rand_seed} to int {rand_seed_new}. This might cause issues.')
-            sample_idx = np.random.RandomState(seed=rand_seed).permutation(len(self))
-        else:
-            sample_idx = np.arange(len(self))
-        train_sample_idx = sample_idx[:n_train_samp]
-        valid_sample_idx = sample_idx[n_train_samp:(n_train_samp+n_valid_samp)]
-        test_sample_idx = sample_idx[(n_train_samp+n_valid_samp):]
-        # create samplers
-        train_sampler = SubsetRandomSampler(train_sample_idx)
-        valid_sampler = SubsetRandomSampler(valid_sample_idx)
-        test_sampler = SubsetRandomSampler(test_sample_idx)
-        # create dataloaders
-        train_loader = DataLoader(self,batch_size=batch_size,sampler=train_sampler,num_workers=num_workers)
-        valid_loader = DataLoader(self,batch_size=batch_size,sampler=valid_sampler,num_workers=num_workers)
-        test_loader = DataLoader(self,batch_size=batch_size,sampler=test_sampler,num_workers=num_workers)
-
-        return train_loader, valid_loader, test_loader
-
-    @property
-    def cummulative_sizes( self ):
-        warnings.warn("cummulative_sizes attribute is renamed to "
-                      "cumulative_sizes", DeprecationWarning, stacklevel=2)
-        return self.cumulative_sizes
-
-def data_transform_normalize( src, trg, scale_factor=1. ):
-    r'''Data transform. Normalizes src, trg pairs through z-scoring.
-    '''
-
-    sample = np.concatenate((src,trg),axis=-1)
-    center = np.mean(sample,axis=-1)
-    std = np.std(sample,axis=-1)
-    src = scale_factor * ((src.T - center)/std).T # is there a better way to align dimensions? einsum?
-    trg = scale_factor * ((trg.T - center)/std).T
-    return (src, trg)
-
 
 def parse_file_info(file_path):
     file_name = os.path.basename(file_path)

From 809374eaace6cdc8126dca195c878a68fc72ec26 Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Tue, 30 Nov 2021 10:31:58 -0800
Subject: [PATCH 6/7] Enh: add docstrings to data.py functions

---
 aopy/data.py | 109 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 108 insertions(+), 1 deletion(-)

diff --git a/aopy/data.py b/aopy/data.py
index 9c519c5e..eeb014cb 100644
--- a/aopy/data.py
+++ b/aopy/data.py
@@ -1044,6 +1044,20 @@ def pkl_read(file_to_read, read_dir):
 
 
 def parse_file_info(file_path):
+    """parse_file_info
+
+    Parses file strings for goose_wireless ECoG and LFP signal data into data parameters.
+
+    Args:
+        file_path (str): path to the file's location
+
+    Returns:
+        exp_file_name (str): JSON experiment data file path
+        mask_file_name (str): binary data mask file path
+        microdrive_name (str): string name of the microdrive type used to collect data in file_path
+        rec_type (str): recording modality reflected in this file ('ECOG', 'LFP', etc.)
+    """
+
     file_name = os.path.basename(file_path)
     data_file_noext = os.path.splitext(file_name)[0]
     data_file_parts = data_file_noext.split('.')
@@ -1057,6 +1071,18 @@ def parse_file_info(file_path):
     return exp_file_name, mask_file_name, microdrive_name, rec_type
 
 def load_experiment_data(exp_file_name):
+    """load_experiment_data
+
+    Reads experiment metadata from an experiment JSON file. Returns the complete data structure as a dictionary and returns electrode data as a pandas DataFrame.
+
+    Args:
+        exp_file_name (str): JSON experiment data file path
+
+    Returns:
+        experiment (dict): dict data object containing experiment metadata. See lab documentation for more information.
+        electrode_df (DataFrame): pandas DataFrame containing microdrive electrode information. Individual channels are indexed along columns, column names are electrode IDs.
+    """
+
     assert os.path.exists(exp_file_name), f'inferred experiment file not found at {exp_file_name}'
     with open(exp_file_name,'r') as f:
         experiment = json.load(f)
@@ -1066,11 +1092,34 @@ def load_experiment_data(exp_file_name):
     return experiment, electrode_df
 
 def load_mask_data(mask_file_name):
+    """load_mask_data
+
+    Loads binary mask data from recording mask files. Binary True values indicate "bad" or noisy data not used in analyses.
+
+    Args:
+        mask_file_name (str): file path to binary mask file
+
+    Returns:
+        mask (numpy.array): numpy array of binary values. Length is equal to the number of time points in the respective data array.
+    """
+
     assert os.path.exists(mask_file_name), f'inferred mask file not found at {mask_file_name}'
     with open(mask_file_name,'r') as f:
-        mask = pkl.load(f)
+        return pkl.load(f)
 
 def read_lfp(file_path,t_range=(0,-1)):
+    """read_lfp
+
+    reads data from a structured binary *lfp file in the goose wireless dataset.
+
+    Args:
+        file_path (str): file path to data file
+        t_range (listlike, optional): Start and stop times to read data. (0, -1) reads the entire file. Defaults to (0,-1).
+
+    Returns:
+        da (numpy.array): numpy array of multichannel recorded neural activity saved in file_path
+        mask (numpy.array): numpy array of binary mask values
+    """
 
     # get local experiment, mask files
     exp_file_name, mask_file_name, microdrive_name, rec_type = parse_file_info(file_path)
@@ -1144,6 +1193,26 @@ def read_lfp(file_path,t_range=(0,-1)):
 
 # wrapper to read and handle clfp ECOG data
 def load_ecog_clfp_data(data_file_name,t_range=(0,-1),exp_file_name=None,mask_file_name=None,compute_mask=True):
+    """load_ecog_clfp_data
+
+    Load ECoG data file from a goose wireless dataset file.
+
+    Args:
+        data_file_name (str): file path to data file
+        t_range (listlike, optional): Start and stop times to read data. (0, -1) reads the entire file. Defaults to (0,-1).
+        exp_file_name (str, optional): File path to experiment data JSON file.
+        mask_file_name (str, optional): File path to data quality mask file. Defaults to None.
+        compute_mask (bool, optional): Compute a data quality mask array if no mask file is given or found. Defaults to True.
+
+    Raises:
+        NameError: If experiment file cannot be found, NameError is raised.
+        NameError: If mask file cannot be found, NameError is raised.
+
+    Returns:
+        data (numpy.array): numpy array of multichannel ECoG data
+        mask (numpy.array): binary mask indicating bad data samples
+        exp (dict): dictionary of experiment data
+    """
 
     # get file path, set ancillary data file names
     exp_file_name, mask_file_name, microdrive_name, rec_type = parse_file_info(data_file_name)
@@ -1231,6 +1300,19 @@ def load_ecog_clfp_data(data_file_name,t_range=(0,-1),exp_file_name=None,mask_fi
 
 # read T seconds of data from the start of the recording:
 def read_from_start(data_file_path,data_type,n_ch,n_read):
+    """read_from_start
+
+    Read data from goose wireless data file. Reads a fixed number of samples from the start of the recording.
+
+    Args:
+        data_file_path (str): file path to data file
+        data_type (numeric type): numpy numeric type reflecting the variable encoding in data_file_path
+        n_ch (int): number of channels saved in data_file_path
+        n_read (int): number of time points to read from data_file_path
+
+    Returns:
+        data (np.array): numpy array of neural recording data saved in data_file_path
+    """
     data_file = open(data_file_path,"rb")
     data = np.fromfile(data_file,dtype=data_type,count=n_read*n_ch)
     data = np.reshape(data,(n_ch,n_read),order='F')
@@ -1240,6 +1322,21 @@ def read_from_start(data_file_path,data_type,n_ch,n_read):
 
 # read some time from a given offset
 def read_from_file(data_file_path,data_type,n_ch,n_read,n_offset,reshape_order='F'):
+    """read_from_file
+
+    Reads recorded neural activity from a goose_wireless file.
+
+    Args:
+        data_file_path (str): file path to data file
+        data_type (numeric type): numpy numeric type reflecting the variable encoding in data_file_path
+        n_ch (int): Number of channels in data_file_path
+        n_read (int): Number of data samples read from data_file_path
+        n_offset (int): Offset point defining where data reading starts
+        reshape_order (str, optional): Data reshaping order. Defaults to 'F'
+
+    Returns:
+        data (np.array): numpy array of neural activity stored in data_file_path
+    """
     data_file = open(data_file_path,"rb")
     if np.version.version >= "1.17": # "offset" field not added until later installations
         data = np.fromfile(data_file,dtype=data_type,count=n_read*n_ch,
@@ -1254,6 +1351,16 @@ def read_from_file(data_file_path,data_type,n_ch,n_read,n_offset,reshape_order='
 
 # read variables from the "experiment.mat" files
 def get_exp_var(exp_data,*args):
+    """get_exp_var
+
+    Generate a list of variable names from a .MAT formatted experiment data
+
+    Args:
+        exp_data (dict): MAT file data dict
+
+    Returns:
+        var_names (list): list of variable names in exp_data
+    """
     out = exp_data.copy()
     for k, var_name in enumerate(args):
         if k > 1:

From 2a125674e38640ed7face8f7b663559bd7fe7c56 Mon Sep 17 00:00:00 2001
From: Michael Nolan <manolan@uw.edu>
Date: Wed, 30 Mar 2022 16:27:04 -0700
Subject: [PATCH 7/7] Add: data transforms to torch.py

---
 aopy/torch.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/aopy/torch.py b/aopy/torch.py
index 05c73676..f3e90660 100644
--- a/aopy/torch.py
+++ b/aopy/torch.py
@@ -65,3 +65,31 @@ def recursive_assign_device(x, device: str):
     else:
         x = x.to(device)
     return x
+
+# transforms - -
+class DropChannels(object):
+    '''
+        Dataset transform to randomly drop channels (i.e. set all values to zero) within a sample.
+        The number of dropped channels is determined by the drop ratio:
+            n_drop = floor(drop_ratio*n_ch)
+        Channel dimension is assumed to be the last indexed tensor dimension. This may need to be
+        adjusted for multidimensional time series data, e.g. spectrograms.
+    '''
+    def __init__(self,drop_ratio=0.1):
+        self.drop_ratio = drop_ratio
+
+    def __call__(self,sample):
+        n_ch = sample.shape[-1]
+        n_ch_drop = floor(self.drop_ratio*n_ch)
+        drop_ch_idx = torch.randperm(n_ch)[:n_ch_drop]
+        sample[:,drop_ch_idx] = 0.
+        return sample
+
+# z-scoring for tensors in pytorch.
+def tensor_zscore(x,dim=0):
+    mean = x.mean(dim=dim).expand([50,-1,-1]).permute(1,0,2)
+    std = x.std(dim=dim).expand([50,-1,-1]).permute(1,0,2)
+    return (x - mean) / std
+
+#-------------------------------------------------------------------
+#-------------------------------------------------------------------
\ No newline at end of file