Search code examples
pythonnumpymatrix-multiplication

Understanding matrix multiplication of ragged nested sequences in NumPy


I am new to NumPy and I am trying to migrate some code I developed in MATLAB for calculating 2x2 transfer functions. Here is the code snippet I wrote.

v = np.arange(-0.5, 0.5, 0.001)
z = np.exp(-1j * 2 * np.pi * v);
Hcoup0 = Hcoup(k0) # those are simple 2x2 matrices
HcoupN = Hcoup(kN)

A1 = np.ones(n_points)
A2 = A1

for n in range(len(k1)):
    A1 = A1 * Hring(z**2, d1[n])
for n in range(len(k2)):
    A2 = A2 * Hring(z**2, d2[n])


# H = np.zeros(n_points, dtype = 'complex_')
# G = H
# for i in range(n_points):
#     Harms = np.array([[A1[i], 0],[0, A2[i]]]) @ np.array([[1, 0], [0, z**1]])

#     HTOT = HcoupN @ Harms @ Hcoup0

#     H[i] = HTOT[0,0]
#     G[i] = HTOT[0,1]
    
Harms = np.array([[A1, 0],[0, A2]], dtype=object) @ np.array([[1, 0], [0, z**1]], dtype=object)
HTOT = HcoupN @ Harms @ Hcoup0

H = HTOT[0,0]
G = HTOT[0,1]

Let me anticipate you that the code does exactly what is supposed to do, so the result is correct, and the matrix multiplication I do is equivalent to what I commented in the code snippet. However, to avoid in the future unpredictable behaviour, I would like to understand a little bit better the logic of the multiplication when the array dimensions are not equal, and the order of operations that it does, because for me in this case it is very convenient that works. I could not find a clear explanation in the web.


Solution

  • The matrix multiplication operator follows the same order of operations that you would anticipate for any ordinary data type. But since the arrays are of dtype=object, it is implemented by invoking their * and + operators.

    Consider your line

    Harms = np.array([[A1, 0],[0, A2]], dtype=object) @ np.array([[1, 0], [0, z**1]], dtype=object)
    

    Let's define a = np.array([[A1, 0],[0, A2]], dtype=object) and b = np.array([[1, 0], [0, z**1]], dtype=object). Both a and b are 2x2 matrices. That the elements of matrices are arrays or scalars, is secondary. From the normal matrix multiplication definition, it follows that

    Harms[0, 0] = a[0, 0] * b[0, 0] + a[0, 1] * b[1, 0]
    Harms[0, 1] = a[0, 0] * b[0, 1] + a[0, 1] * b[1, 1]
    Harms[1, 0] = a[1, 0] * b[0, 0] + a[1, 1] * b[1, 0]
    Harms[1, 1] = a[1, 0] * b[0, 1] + a[1, 1] * b[1, 1]
    

    So, for example the first line is equivalent to

    Harms[0, 0] = A1 * 1 + 0 * 0
    

    At this point the normal broadcasting and type conversion rules of numpy apply.

    If you change the shape of those matrices, the operations will expand as you would expect from any naive matrix multiplication algorithm.