In numpy, multipy two structured matrices concisely

Question:

I have two matrices. The first has the following structure:

[[1, 0, a],
 [0, 1, b],
 [1, 0, c],
 [0, 1, d]]

where 1, 0, a, b, c, and d are scalars. The matrix is 4 by 3

The second is just a 2 by 3 matrix:

[[r1],
[r2]]

where r1 and r2 are the first and second rows respectively, each having 3 elements.
I would like the output to be:

[[r1, 0, a*r1],
 [0, r1, b*r1],
 [r2, 0, c*r2],
 [0, r2, d*r2]]

which would be a 4 by 9 matrix.
This is similar to the Kronecker product, except separately for each row of the second matrix. Of course this could be done with cumbersome loops which I want to avoid.
How can I do this concisely?

Asked By: Yakov Dan

||

Answers:

You can do exactly what you said in the last line: do a separate Kronecker product for each row of the second column and then concatenate the results.

Let’s assume that the two matrices are called x (4 by 3) and y (2 by 3). The first thing to do is to split x in two parts because only half matrix participates in each part of the product.

x = x.reshape(2, 2, 3)

Then you can calculate the two products separately:

z0 = np.kron(x[0], y[0])
z1 = np.kron(x[1], y[1])

Finally, concatenate the two results along the first axis:

z = np.concatenate([z0, z1], axis=0)

Or if, like me, you enjoy big ugly one-liners you can do:

z = np.concatenate([np.kron(xr, yr) for xr, yr in zip(x.reshape(2, 2, 3), y)], axis=0)

In the general case you mentioned in the comments, it would become:

z = np.concatenate([np.kron(xr, yr) for xr, yr in zip(x.reshape(int(n / 2), 2, 3), y)], axis=0)

This gives equal results to the explicit loop, which can be numba.jit compiled I believe:

def solve_explicit(x, y):
    # sanity checks
    assert x.shape[0] == 2*y.shape[0]
    assert x.shape[1] == y.shape[1]

    n = x.shape[0]
    z = np.zeros((n, 9))
    for i in range(n):
        for j in range(3):
            for k in range(3):
                z[i, k + 3 * j] = x[i, j] * y[int(i / 2), k]
    return z
Answered By: Gianluca Micchi

Using broadcasting, with x.shape (n, 3), and y.shape (n//2, 3):

out = (x.reshape(-1, 2, 3, 1) * y.reshape(-1, 1, 1, 3)).reshape(-1, 9)
Answered By: Chrysophylaxs

I personally would use np.einsum in this situation because I think it’s easier to understand than broadcasting.

import numpy as np

(a, b, c, d) = np.random.rand(4)
x = np.array([[1, 0, a], [0, 1, b], [1, 0, c], [0, 1, d]])
y = np.random.rand(2, 3)
z = np.einsum("ij,ik->ijk", x.reshape(-1, 6), y).reshape(-1, 9)

# timeit magic commands.
# %timeit -n 50000 np.einsum("ij,ik->ijk", x.reshape(-1, 6), y).reshape(-1, 9)
# %timeit -n 50000 (x.reshape(-1, 2, 3, 1) * y.reshape(-1, 1, 1, 3)).reshape(-1, 9)

Some good references on Einstein summation in NumPy: [2, 3, 4].

Answered By: airalcorn2