python - How to replace one dimension of an n-dimensional array with another?

Question

I have a numpy array of shape (4, 30, 13, 7000). This is experimental data. The first three dimensions represent experimental conditions. The last dimension represents 7000 ms for each combination of the first three dimensions. So there are 1560 7000ms lists in the array.

I've built a sliding-window mean function which I perform on every 7000ms list:

def windowed_mean(4D_list)

    for trial in 4D_list:
        for neuron in trial:
            for timebin in neuron:
                chunk = timebin #chunk equals every 7000ms timespan. 
                window_size = 250 #ms
                i = 0
                while i < len(chunk) - window_size + 1: #sliding window average
                    window = chunk[i : i + window_size] #generates window
                    window_average = sum(window) / window_size #takes window average
                    moving_average.append(window_average) #appends window average to #moving_average list 
                    i += 25 #step size
    
                stored_averages.append(window_average)
                moving_average.clear()

    
    print(len(stored_averages)) #this list contains the stored windowed averages in order

My question is, how do I replace the original array's fourth dimension (time) with the new windowed means stored in stored_averages? I'd ideally like to end up with a new array of shape 4, 30, 13, 271, 271 because that's the number of windowed means I get per 7000 ms trial.

Answer 1

If you want a moving average, take a look at scipy.ndimage.convolve1d. A sliding window is just convolution with a box function of the appropriate width and height.

You are looking for something like

def windowed_mean(arr, n, axis=-1):
    box = np.full(n, 1.0 / n)
    return ndimage.convolve1d(arr, box, axis)

This will return an array of the same size as the original. You likely want something that does not include the partially convolved elements, so you can trim (n - 1) // 2 from the left and n // 2 from the right. Using integer division like that ensures that the trim is correct for both even and odd windows:

    return  ndimage.convolve1d(arr, box, axis)[..., (n - 1) // 2:-(n // 2)]

You can do the same convolution using a 1D convolver like np.convolve. This would require your data to be arranged so that the dimension over which you are convolving is contiguous. This is likely the case here, since numpy uses C order by default:

def windowed_mean(arr, n):
    box = np.full(n, 1.0 / n)
    conv = np.convolve(arr.ravel(), box)
    return conv[n - 1:].reshape(arr.shape)[..., :1 - n]

To operate over not-the-last dimension, you would have to move the axis of interest to the end. Keep in mind that ravel will make a copy of the data in that case:

def windowed_mean(arr, n, axis=-1):
    box = np.full(n, 1.0 / n)
    a = np.moveaxis(arr, axis, -1).ravel()
    conv = np.convolve(a, box)
    conv = conv[n - 1:].reshape(arr.shape)[..., :1 - n]
    return np.moveaxis(conv, -1, axis)