Functions for faster data augmentation of timeseries

faster_splits.py

def split_Nd_array(array: np.ndarray, nsplits: int) -> List[np.ndarray]:
    if array.ndim == 1:
        indices = range(0, len(array) - 31, nsplits)
        segments = [np.take(array, np.arange(i, i + 32), axis=0).copy() for i in indices]
    else:
        indices = range(0, array.shape[1] - 31, nsplits)
        segments = [np.take(array, np.arange(i, i + 32), axis=1).copy() for i in indices]
    return segments

def create_new_segments_from_splits(segment: _Segment, nsplits: int) -> List[_Segment]:
    new_segments = []
    if segment['output_array'].shape[1] > 32:
        output_array = split_Nd_array(array=segment['output_array'], nsplits=nsplits)
        bursts = split_Nd_array(array=segment['doppler_burst'], nsplits=nsplits)
        new_segments.extend([_Segment(segment_id=f'{segment["segment_id"]}_{j}',
                                      output_array=array,
                                      doppler_burst=bursts[j],
                                      target_type=segment['target_type'],
                                      segment_count=1)
                             for j, array in enumerate(output_array)])

    else:
        new_segments.append(segment)
    return new_segments