How to directly vectorize a function that takes indices as an input and returns a 2x2 np.array

Question

Let's say I have a function f(i,j,a) that takes the row and column from a 2D-numPy array a and returns another 2 x 2 numPy array with i,j at the top left corner bi.e

def f(i,j,a):
 b = a[np.arange(2)[:, np.newaxis]  + i, np.arange(2) + j]
 return b

(The above code could be generalized by replacing the 2 with some other number). I want to vectorize this process by defining f(i_array,j_array,a) by essentially iterating over the i,j. This is simple with list comprehension:

b_array = np.array([ f(i,j,a) for i in i_array, for j in j_array]).

Is there a more efficient way to do this directly with numPy functions/indexing or is list comprehension the best I can do? I know that np.vectorize does this but I don't believe it is any faster than a simple list comprehension.

Answer 1

I figured it out (but my solution is a bit janky so any revisions would be appreciated). My code looks like:

import numpy as np

def f(i_array,j_array, a):
 i_transposed, j_transposed = i_array.T, j_array.T
 i_matrix = np.arange(2)[:, np.newaxis][:, np.newaxis]+ i_transposed
 j_matrix = np.arange(2)[:, np.newaxis]+ j_transposed
 b_grids = np.swapaxes(grid[i_matrix, j_matrix].T, 1,2)

 return b_grids

Here, we take advantage of broadcasting to generate a np.array of the previous matrix. The i_matrix term takes the np.arange(2), increases its dimension by 2 and uses broadcasting to make a (2, 1, 2) array with all the possible i_transposed added to 2. The same happens with j_matrix except it is 1 dimension lower. Then we use numPy indexing to broadcast in grid[i_matrix, j_matrix]. The dimensions are wrong because of the way we broadcasted and so we have to use np.swapaxes.

This should work somewhat more generally by replacing np.arange(2) with whatever offset you want your coordinates to have.