Python filtering signals in a 4D structure

Question

I have a 4D structure containing EEG data with shape (3432,1,30,512)

This represents 3432 trials of data, where each trial contains 512 time samples for 30 electrodes

For each trial, I want to filter each electrode with certain filter parameters, which results in a new 2D (30,512) filtered array. Then, I want to stack these along axis 1 of the 4D structure.

Currently, I'm iterating over each trial, then each electrode, filtering the signal and inserting everything back into the 4D:

from scipy import signal    

NUM_CHANS = 30
NUM_TIMESTAMPS = 512
FREQ_BANDS = ((0.1, 3), (4, 7), (8, 13), (16, 31))

all_data = np.reshape(all_data, (-1, 1, NUM_CHANS, NUM_TIMESTAMPS))  # (3432,1,30,512)

for i in range(all_data.shape[0]):
    num_band = 1
    for band in FREQ_BANDS:
        lower = float(band[0])/(SAMPLE_RATE/2)
        upper = float(band[1])/(SAMPLE_RATE/2)

        # Design new filter for the current frequency band
        b, a = signal.butter(2, [lower, upper], 'bandpass')

        temp_trial = np.zeros((NUM_CHANS, NUM_TIMESTAMPS))

        for ch in range(NUM_CHANS):
            # Filter the current electrode
            output_signal = signal.filtfilt(b, a, all_data[i,0,ch,:])
            temp_trial[ch,:] = output_signal

        # Insert temp_trial (2D) into all_data (4D) along axis 1

        num_band += 1

Iterating over trials and electrodes is extremely slow (takes about 2 hours to complete the entire loop). Is there a more efficient way to apply this filter over all electrodes/trials?

I've been trying to find a way apply the filter over 2D, so I don't need to iterate over electrodes, but haven't been able to find anything.

EDIT:

Would this be the correct way to use filtfilt axis argument?

all_data = np.reshape(all_data, (-1, 1, NUM_CHANS, NUM_TIMESTAMPS))

filtered_data = [all_data]

for band in FREQ_BANDS:
    lower = float(band[0])/(SAMPLE_RATE/2)
    upper = float(band[1])/(SAMPLE_RATE/2)
    b, a = signal.butter(2, [lower, upper], 'bandpass')

    output_signal = signal.filtfilt(b, a, all_data, axis=3)

    filtered_data.append(output_signal)

all_data = np.concatenate(filtered_data, axis=1)

Answer 1

Here's one change that might improve performance. filtfilt has an axis argument, which allows you to apply the same 1-d filter along an axis of an n-dimensional array. You can replace this code

    temp_trial = np.zeros((NUM_CHANS, NUM_TIMESTAMPS))
    for ch in range(NUM_CHANS):
        # Filter the current electrode
        output_signal = signal.filtfilt(b, a, all_data[i,0,ch,:])
        temp_trial[ch,:] = output_signal

with

    temp_trial = signal.filtfilt(b, a, all_data[i,0,:,:], axis=1)

And since the default axis is axis=-1, you can leave out the axis argument if you prefer:

    temp_trial = signal.filtfilt(b, a, all_data[i,0,:,:])

---

You can do even better, by rearranging the outer loops. Make the "band" loop the outer loop. Then, for each band, you can apply filtfilt once to the 3-d array all_data[:, 0, :, :].

Something like this:

shp = all_data.shape

# This assumes `all_data` has shape (3432,1,30,512), but I don't
# think you need that trivial extra dimension in there if you
# preallocate `filtered_data` like I am doing here.
filtered_data = np.empty(shp[0:1] + (len(FREQ_BANDS),) + shp[2:])

for k, band in enumerate(FREQ_BANDS):
    lower = float(band[0])/(SAMPLE_RATE/2)
    upper = float(band[1])/(SAMPLE_RATE/2)

    # Design new filter for the current frequency band
    b, a = signal.butter(2, [lower, upper], 'bandpass')

    filtered_data[:, k, :, :] = filtfilt(b, a, all_data[:,0,:,:])

This doesn't include the original all_data in filtered_data. If you want that, use len(FREQ_BANDS)+1 in the call to np.empty() that creates filtered_data, set filtered_data[:,0,:,:] = all_data, and use k+1 instead of k in assignment statement that saves the result of the filtfilt() call.