How to vectorize/tensorize operations in numpy with irregular array shapes

Viewed 1094

I would like to perform the operation

form1

If form2 had a regular shape, then I could use np.einsum, I believe the syntax would be

np.einsum('ijp,ipk->ijk',X, alpha)

Unfortunately, my data X has a non regular structure on the 1st (if we zero index) axis.

To give a little more context, form3 refers to the p^th feature of the j^th member of the i^th group. Because groups have different sizes, effectively, it is a list of lists of different lengths, of lists of the same length.

form4 has a regular structure and thus can be saved as a standard numpy array (it comes in 1-dimensional and then I use alpha.reshape(a,b,c) where a,b,c are problem specific integers)

I would like to avoid storing X as a list of lists of lists or a list of np.arrays of different dimensions and writing something like

A = []
for i in range(num_groups):
    temp = np.empty(group_sizes[i], dtype=float)
    for j in range(group_sizes[i]):
        temp[i] = np.einsum('p,pk->k',X[i][j], alpha[i,:,:])
    A.append(temp)

Is this some nice numpy function/data structure for doing this or am I going to have to compromise with some only partially vectorised implementation?

1 Answers
Related