Search code examples
pandasnumpymatrix-multiplicationnumpy-einsum

Python: multiply 2 time series of matrices


I have 2 time-series of matrices that I would like to multiply together. I have the data stored as a pandas MultiIndex frame. Both frames share the same first axis, which are k dates. So, the first matrix A has k dates, and for each date it has a n x n matrix, making it k x n x n. The second matrix B also has k dates and for each date it has an n by m matrix, making it k x n x m. For each date k, I would like to multiple the n x n part of A with the n x m part of B. I have changed my MultiIndex into a numpy array in a bid to speed up the operation as k is circa 15,000,000. I have set this up as: AB = np.array([np.dot(A.loc[date_], B.loc[date_]) for date_ in B.index.levels[0]]). Unfortunately, this takes many, many hours to run. I have heard that using np.einsum is much, much quicker, but I have thusfar been unable to get it to work. Does anyone have any suggestions on how to speed this up and if possible, use np.einsum to achieve that same result in less time?


Solution

  • After playing around with einsum for hours, I finally managed to solve this: np.einsum('ijk,ikl->ijl', A, B)