broadcasting numpy kronecker product

Question

I have to calculate the Kronecker's product in the following loop:

import numpy as np
M=48
G=120
H=5
P=6
we=np.random.normal(size=(P,P,G))
qu=np.random.normal(size=(P,G))
f = np.random.normal(size=(H,G))
P, G = qu.shape
nu = np.full((P * H, 1), 0.0)
de = np.full((P * H, P * H), 0.0)
for g in range(G):
    qu_g = qu[:, g].reshape((-1, 1))
    f_g = f[:, g].reshape((-1, 1))
    we_g = we[:, :, g]
    nu += (qu_gt[:, None]*f_g[None, :]).reshape(P*H,1)
    de += (we_g[:, None, :, None]*f_g.dot(f_g.T)[None, :, None, :]).reshape(P*H,P*H)
nu = nu.sum(axis=0)
de = de.sum(axis=0)
ga2 = np.linalg.solve(de, nu).reshape((P, H))

Notice that (qu_g[:, None]*f_g[None, :]).reshape(P*H,1) is equivalent to np.kron(qu_g, f_g) but faster thanks to broadcasting. Analgously for (we_g[:, None, :, None]*f_g.dot(f_g.T)[None, :, None, :]).reshape(P*H,P*H).

I would like to exploit numpy broadcasting to avoid the loop and speed up the process. What I mean is that my aim is to do everything at once exploiting the fact that, given a 3d array, numpy carries out the dot product to the last two dimensions using matmul, such that one can avoid looping over the first dimension. I would like to exploit the same logic also for kron.

However, I cannot obtain the same results. My attempt for the nu component: (qu[:, None].transpose(2,0,1)*f[:, None].transpose(2,1,0)).reshape(G, P*H, 1) does not give the same result. How can I proceed here?

Answer 1

Here's the solution I propose:

sf = np.matmul(f[:,None].transpose(2,0,1),f[None, :].transpose(2,0,1))
nu = np.einsum('ai,ak->aik',qu.T,f.T).reshape(qu.shape[1],qu.shape[0]*f.shape[0], 1)
de = np.einsum('aij,akl->aikjl',we.transpose(2,0,1),sf).reshape(we.transpose(2,0,1).shape[0],we.transpose(2,0,1).shape[1]*sf.shape[1],we.transpose(2,0,1).shape[2]*sf.shape[2])
nu = nu.sum(axis=0)
de = de.sum(axis=0)
np.linalg.solve(de, nu).reshape((P, H))