Search code examples
pythonnumpysubmatrix

Extract sub rows with varying sizes from a big 2D NumPy Array


I have a NumPy Array with size say 3*10, I would like to extract sub rows with varying sizes from each row. The sub rows are centered in the middle pixel with varying pixel sizes. Then I take the average number of each subrow. I have a pseudo example below:

import numpy as np
arr = np.arange(1,31).reshape((3,10))
pixel_size = np.array([2,3,1])
## the subrow centers in the middle of the array, index 5
mask = [[5-2:5+2],[5-3:5+3],[5-1:5+1]] ## index for each row
### submatrix = arr[;,mask]
submatrix = [[3,4,5,6],[12,13,14,15,16,17],[24,25]]
## output = np.mean(submatrix, axis=1) output is the average number of each row in the submatrix
output = [4.5,14.5,24.5]

If I have over 10 millions of rows, how can I handle this situation fast.


Solution

  • You can do it using list comprehensions and index slicing:

    import numpy as np
    
    arr = np.arange(1,31).reshape((3,10))
    pixel_size = np.array([2,3,1])
    
    middle_ind = int(arr.shape[1]/2.)
    print middle_ind
    sub_arr = [arr[i,middle_ind - pixel_size[i]:middle_ind + pixel_size[i]] for i in range(len(pixel_size))]
    print('sub_arr: ', sub_arr)
    output = [np.mean(item) for item in sub_arr]
    print('output: ', output)
    
    > sub_arr: [array([4, 5, 6, 7]), array([13, 14, 15, 16, 17, 18]), array([25, 26])]
    > output: [5.5, 15.5, 25.5]
    

    Your submatrix is a list not an array so it's more difficult to vectorize operations. You might want to think about restructuring your code to take advantage of matrix operations.