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

Question

I would like to perform the operation

If 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, 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.

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?

Answer 1

I know this sounds obvious, but, if you can afford the memory, I'd start just by checking the performance you get simply by padding the data to have a uniform size, that is, simply adding zeros and perform the operation. Sometimes a simpler solution is faster than a more supposedly optimal one that has more Python/C roundtrips.

If that doesn't work, then your best bet, as Tom Wyllie suggested, is probably a bucketing strategy. Assuming X is your list of lists of lists and alpha is an array, you can start by collecting the sizes of the second index (maybe you already have this):

X_sizes = np.array([len(x_i) for x_i in X])

And sort them:

idx_sort = np.argsort(X_sizes)
X_sizes_sorted = X_sizes[idx_sort]

Then you choose a number of buckets, which is the number of divisions of your work. Let's say you pick BUCKETS = 4. You just need to divide the data so that more or less each piece is the same size:

sizes_cumsum = np.cumsum(X_sizes_sorted)
total = sizes_cumsum[-1]
bucket_idx = []
for i in range(BUCKETS):
    low = np.round(i * total / float(BUCKETS))
    high = np.round((i + 1) * total / float(BUCKETS))
    m = sizes_cumsum >= low & sizes_cumsum < high
    idx = np.where(m),
    # Make relative to X, not idx_sort
    idx = idx_sort[idx]
    bucket_idx.append(idx)

And then you make the computation for each bucket:

bucket_results = []
for idx in bucket_idx:
    # The last index in the bucket will be the biggest
    bucket_size = X_sizes[idx[-1]]
    # Fill bucket array
    X_bucket = np.zeros((len(X), bucket_size, len(X[0][0])), dtype=X.dtype)
    for i, X_i in enumerate(idx):
        X_bucket[i, :X_sizes[X_i]] = X[X_i]
    # Compute
    res = np.einsum('ijp,ipk->ijk',X, alpha[:, :bucket_size, :])
    bucket_results.append(res)

Filling the array X_bucket will probably be slow in this part. Again, if you can afford the memory, it would be more efficient to have X in a single padded array and then just slice X[idx, :bucket_size, :].

Finally, you can put back your results into a list:

result = [None] * len(X)
for res, idx in zip(bucket_results, bucket_idx):
    for r, X_i in zip(res, idx):
        result[X_i] = res[:X_sizes[X_i]]

Sorry I'm not giving a proper function, but I'm not sure how exactly is your input or expected output so I just put the pieces and you can use them as you see fit.