Python3: Vectorizing Nested Loops

I have this function: def fun(x): # x is a vector with size: (size_x*size_y) = n c = 0 f_vec = np.zeros((size_x*size_y)) for i in range(size_x): for j in r

Solution 1:

You can use broadcasting for this:

(
    x.reshape(size_x, size_y) *
    np.arange(size_x)[:, None] *
    np.arange(size_y)
).ravel()

or Einstein Summation form

np.einsum(
    'ij,i,j->ij',
    x.reshape(size_x, size_y),
    np.arange(size_x),
    np.arange(size_y)
).ravel()

Solution 2:

Essentially, this is the same as Nils Werner's answer, but I find it easier to understand the i*j part as the 2D ndarray np.outer(np.arange(x_size), np.arange(y_size), and then do broadcasting:

(x.reshape(x_size, y_size) * np.outer(np.arange(x_size), np.arange(y_size)).ravel()

You can pre-calculate the i*j part if you're doing this repeatedly for the same values of x_size and y_size.