Performing a special multiplication on two square matrices

Question

I have a function called mult2 that takes in two square NumPy matrices and returns a matrix that is the result of a special multiplication:

p=P(1,344)
PI,PIT=[],[]
for i in xrange(344):
    PI.append(p**i)
    PIT.append((p**i).transpose())

def mult2(a,b):
    C=np.sum((a*b)**2,axis=1)
    CE1=np.matrix([[0.0] for _ in xrange(len(a))])
    API,PITB=[],[]
    AB=a*b
    for i in xrange(len(a)):
        API.append(np.diag(np.diag(a*PI[i])))
        PITB.append(np.diag(np.diag(PIT[i]*b)))
        CE1+=np.sum(API[i]*PITB[i]*AB,axis=1)
    CE2=np.matrix([[0.0] for _ in xrange(len(a))])
    BPI,PITA=[],[]
    for i in xrange(len(a)):
        BPI.append(np.diag(np.diag(b*PI[i])))
        PITA.append(np.diag(np.diag(PIT[i]*a)))
        CE2+=np.sum(a*BPI[i]*PITA[i]*b,axis=1)
    CE3=np.matrix([[0.0] for _ in xrange(len(a))])
    for i in xrange(len(a)):
        CE3+=np.sum(AB*API[i]*PITB[i],axis=1)
    CE1E2=np.matrix([[0.0] for _ in xrange(len(a))])
    for i in xrange(len(a)):
        CE1E2+=np.sum(API[i]*PITB[i]*API[i]*PIT[i]*b,axis=1)
    CE2E3=np.matrix([[0.0] for _ in xrange(len(a))])
    for i in xrange(len(a)):
        CE2E3+=np.sum(a*BPI[i]*PITA[i]*BPI[i]*PIT[i],axis=1)
    CE1E3=np.matrix([[0.0] for _ in xrange(len(a))])
    for i in xrange(len(a)):
            for j in xrange(len(a)):
                CE1E3+=np.sum(np.matrix(API[i]*PITB[i]*API[j]*PITB[j]),axis=1)
    CE1E2E3=np.matrix([[0.0] for _ in xrange(len(a))])
    for i in xrange(len(a)):
        CE1E2E3+=np.sum(np.matrix((API[i]*PITB[i])**2),axis=1)
    return C-(CE1+CE2+CE3-CE1E2-CE1E3-CE2E3+CE1E2E3)

p is another matrix of the same dimensions (n by n) as a and b specially a circulant matrix.

My trouble is that these matrices are 344 by 344 and I am looking to run this function lots of times.

Storing the matrices that are used later in the function in the lists actually sped up the computation significantly.

I know this is a very intensive computation so it should take a while to run. That being said, can anyone spot any improvements I can make?

Answer 1

From what I can tell, you could combine all of the for loops into only 2 distinct loops:

1. A loop to setup the matrices used for later calculations
2. A loop to do the rest of the calculations

Instead of performing list comprehension to create an array to pass to each numpy matrix, you could simply multiply a base array by the len(a) to create said array:

base_list = [[0.0]]*len(a)  # Creates [[0.0], [0.0], [0.0], ... ]

This way you do not temporarily store each element in the xrange object without needing to.

Use more descriptive variable names. What is a, b? Using more descriptive variable names helps improve comprehension and readability.

---

Here is my version of mult2():

def mult2(a, b):

    AB = a*b
    API, BPI, PITA, PITB = [],[],[],[]

    # Base list used to seed each np.matrix()
    base_list= [[0.0]]*len(a)

    C = np.sum((a*b)**2, axis=1)
    CE1 = np.matrix(base_list)
    CE2 = np.matrix(base_list)
    CE3 = np.matrix(base_list)
    CE1E2 = np.matrix(base_list)
    CE2E3 = np.matrix(base_list)
    CE1E3 = np.matrix(base_list)
    CE1E2E3 = np.matrix(base_list)

    # Create arrays used in later matrix calculations
    for index in xrange(len(a)):
        API.append(np.diag(np.diag(a*PI[index])))
        PITB.append(np.diag(np.diag(PIT[index]*b)))
        BPI.append(np.diag(np.diag(b*PI[index])))
        PITA.append(np.diag(np.diag(PIT[index]*a)))

    # Do matrix calculations
    for index in xrange(len(a)):
        CE1 += np.sum(API[index]*PITB[index]*AB, axis=1)
        CE2 += np.sum(a*BPI[index]*PITA[index]*b, axis=1)
        CE3 += np.sum(AB*API[index]*PITB[index], axis=1)
        CE1E2 += np.sum(API[index]*PITB[index]*API[index]*PIT[index]*b, axis=1)
        CE2E3 += np.sum(a*BPI[index]*PITA[index]*BPI[index]*PIT[index], axis=1)
        CE1E2E3 += np.sum(np.matrix((API[index]*PITB[index])**2), axis=1)

        for pos in xrange(len(a)):
            CE1E3 += np.sum(
              np.matrix(API[index]*PITB[index]*API[pos]*PITB[pos]), axis=1)

    return C - (CE1 + CE2 + CE3 - CE1E2 - CE1E3 - CE2E3 + CE1E2E3)

---

Remember, good code is readable code: group like statements together.

Also, PEP8 is the Python standard style convention. It provides a very nice style guide for code and helps code look as Pythonic as possible.