Optimal multiplication of two 3D arrays having a variable dimension

I would like to multiply tensors R = {R_1, R_2, ..., R_M} and X = {X_1, X_2, ..., X_M} where R_i and X_i are 3×3 and 3×N_i matrices, respectively. How can I make maximum use of NumPy functionalities during the formation of the R_i × X_i arrays?

My MWE is the following:

import numpy as np

np.random.seed(0)

M = 5
R = [np.random.rand(3, 3) for _ in range(M)]
X = []
for i in range(M):
    N_i = np.random.randint(1, 6)
    X_i = np.random.rand(3, N_i)
    X.append(X_i)
    
result = np.zeros((3, 0))
for i in range(M):
    R_i = R[i]
    X_i = X[i]
    result = np.hstack((result, np.dot(R_i, X_i)))

print(result)

>Solution :

I don’t think there is much more to do.
Depending on how the N_i are distributed, for example, if half of N_i are 5, and other half 10, it might help to have two X (one M×5 X for the 5 1st columns, and one (M/2)×5 X for the others). That is, generally speaking, to work column-wise rather than "X-wise"
Or, in some other cases, to fill X with 0s (for example, if N_i is 10000 for 99% of the i, and 9999 for the remaining %).

But, roughly, numpy is to handle homogeneous dimensions array.

One point tho: never ever use hstack stack vstack, append or things like that in a for loop. That is very slow. That implies reallocating a new array and copying all the content of the former array into the new one.

Either use plain simple python lists to accumulate the results. And do the stacking at the end only

res=[]
for i in range(M):
    R_i = R[i]
    X_i = X[i]
    res.append(np.dot(R_i, X_i))

np.hstack(res)

Or allocate in advance the the result numpy array, and store the result as you iterate.

result=np.empty((3,sum(x.shape[1] for x in X)))
k=0
for i in range(M):
    R_i = R[i]
    X_i = X[i]
    result[:,k:k+X_i.shape[1]]=np.dot(R_i,X_i)
    k+=X_i.shape[1]