Produit Kronecker efficace avec matrice d'identité et matrice régulière - NumPy/Python

Question

Je travaille sur un projet python et utilise numpy. J'ai souvent besoin de calculer les produits de matrices de Kronecker par la matrice d'identité. C'est un gros goulot d'étranglement dans mon code, donc je voudrais les optimiser. Il y a deux types de produits que je dois prendre. La première est:

np.kron(np.eye(N), A) 

Celui-ci est assez facile d'optimiser en utilisant simplement scipy.linalg.block_diag. Le produit est équivalent à:

la.block_diag(*[A]*N) 

Qui est environ 10 fois plus rapide. Cependant, je ne suis pas sûr sur la façon d'optimiser le deuxième type de produit:

np.kron(A, np.eye(N)) 

Y a-t-il un tour similaire que je peux utiliser?

Answer 1

Une approche consiste à initialiser un tableau de sortie de 4D, puis d'y affecter des valeurs de A. Une telle affectation permettrait de diffuser des valeurs et c'est là que nous aurions l'efficacité en NumPy.

Ainsi, une solution serait comme si -

# Get shape of A 
m,n = A.shape 

# Initialize output array as 4D 
out = np.zeros((m,N,n,N)) 

# Get range array for indexing into the second and fourth axes 
r = np.arange(N) 

# Index into the second and fourth axes and selecting all elements along 
# the rest to assign values from A. The values are broadcasted. 
out[:,r,:,r] = A 

# Finally reshape back to 2D 
out.shape = (m*N,n*N) 

Mettez en fonction -

def kron_A_N(A, N): # Simulates np.kron(A, np.eye(N)) 
    m,n = A.shape 
    out = np.zeros((m,N,n,N),dtype=A.dtype) 
    r = np.arange(N) 
    out[:,r,:,r] = A 
    out.shape = (m*N,n*N) 
    return out 

Pour simuler np.kron(np.eye(N), A), swap simplement les opérations le long de la première et la deuxième et de même pour la troisième et quatrième axes -

def kron_N_A(A, N): # Simulates np.kron(np.eye(N), A) 
    m,n = A.shape 
    out = np.zeros((N,m,N,n),dtype=A.dtype) 
    r = np.arange(N) 
    out[r,:,r,:] = A 
    out.shape = (m*N,n*N) 
    return out 

Timings -

In [174]: N = 100 
    ...: A = np.random.rand(100,100) 
    ...: 

In [175]: np.allclose(np.kron(A, np.eye(N)), kron_A_N(A,N)) 
Out[175]: True 

In [176]: %timeit np.kron(A, np.eye(N)) 
1 loops, best of 3: 458 ms per loop 

In [177]: %timeit kron_A_N(A, N) 
10 loops, best of 3: 58.4 ms per loop 

In [178]: 458/58.4 
Out[178]: 7.842465753424658