0
$\begingroup$

I need to compute a series of matrix multiplications involving 3x3 matrices and a series of matrix norms also involving 3x3 matrices and I wonder how I can set these computations up with numpy such that they are computed quickly (preferably below 10ms).

I am given matrices $A,B \in \mathbb{C}^{3 \times 3}$ and I need to compute the following:

$$ \tilde{R}_i=AC_iB $$

$$ o_i=||\tilde{R}\tilde{R}^{T}-I_3||_F $$

Where $\mathbb{C}_i \in \mathbb{C}^{3 \times 3}, i\in[1,...,N]$ are matrices that are known beforehand and do not change.

My current approach looks like this:

import numpy as np
from numpy.linalg import norm
import time

#some constants
N=2**14
C=np.random.rand(3,3,N)+1j*np.random.rand(3,3,N)

#function that does computations
def compute_o(A,B):
    global C
    #perform the computations
    ####################################
    T=np.einsum("ij,jkl",A,C)
    R=np.einsum("lkn,kj",T,B)
    ####################################

#generate random inputs
A=np.random.rand(3,3)+1j*np.random.rand(3,3)
B=np.random.rand(3,3)+1j*np.random.rand(3,3)

#perform computations and estimate time
tic=time.time()
compute_o(A,B)
toc=time.time()
print("Time ellapsed: %fms"%(1e3*(toc-tic)))

Basically, I did the initial batch of matrix multiplications using einsum which on my computer takes around 2ms. However, I don't quite understand how I can compute the second step without using some sort of loops. Does anyone have an idea?

$\endgroup$

1 Answer 1

2
$\begingroup$

You can use einsum with explicit output indices, running over dimension n without summing.

T = np.einsum("ij,jkn", A, C)
R = np.einsum("ijn,jk", T, B)
S = np.einsum("ijn,kjn->ikn", R, R) - np.eye(3).reshape((3,3,1))
o = np.linalg.norm(S, ord="fro", axis=(0,1))
$\endgroup$
1
  • 2
    $\begingroup$ You could also use S=np.matmul(R,R,axes=[(0,1),(1,0),(0,1)]) - np.eye(3).reshape((3,3,1)) instead of the einsum, but surprisingly this turned out to perform slightly worse for me. No idea why, though I suspect somewhere there's a different loop order which isn't as cache friendly since the results do differ at floating point precision. $\endgroup$ Apr 23, 2023 at 15:23

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.