EtienneCmb / visbrain

A multi-purpose GPU-accelerated open-source suite for brain data visualization
http://visbrain.org
Other
241 stars 64 forks source link

[ENH] Speed up `_index_to_events` and `_events_to_index` #104

Open snwnde opened 2 years ago

snwnde commented 2 years ago

I opened https://github.com/raphaelvallat/yasa/issues/50#issue and the same improvement could be done here, too. The idea is to avoid for loop, which takes more time than numpy ufuncs, especially when the number of events is huge.

My proposal for _index_to_events is the following:

def _index_to_events(x):
    x_copy = np.copy(x)
    x_copy[:,1] += 1 
    split_idx = x_copy.reshape(-1).astype(int)
    full_idx = np.arange(split_idx.max())
    index = np.split(full_idx, split_idx)[1::2]
    index = np.concatenate(index)
    return index

This is fully compatible with the current implementation according to my tests. For benchmarks, see https://github.com/raphaelvallat/yasa/issues/50#issuecomment-1000740304.

For _events_to_index I tend to use scipy,

def _events_to_index(x):
    from scipy.ndimage.morphology import distance_transform_edt
    full_idx = np.arange(x.astype(int).max() + 3)
    tool_arr = np.zeros_like(full_idx)
    tool_arr[x+1] = 1
    distance_arr = distance_transform_edt(tool_arr)
    edge_bool = (distance_arr == 1)
    edge_idx = full_idx[edge_bool] - 1
    start_idx = edge_idx[0::2]
    end_idx = edge_idx[1::2]
    return np.vstack((start_idx, end_idx)).T

Two potential problems here: first, there might be a reason to avoid scipy here; second, this is not fully compatible with the current implementation, as it would merge events that overlap in time. Are they really problems or not in the scope of the usage of this function?

Please tell me how you think about it. Is it really beneficial to do so, etc :)