En mi proyecto actual necesito "retorcerse"Dos matrices tridimensionales de una manera ligeramente inusual:

Supongamos que tenemos dos matrices tridimensionales A y B con las dimensiones dimA y dimB (iguales para cada eje). Ahora queremos crear una tercera matriz C con las dimensiones dimA + dimB para cada eje.

Las entradas de C se calculan como:

c_{x1+x2,y1+y2,z1+z2} += a_{x1,y1,z1} * b_{x2,y2,z2}

Mi versión actual es sencilla:

dimA = A.shape[0]
dimB = B.shape[0]
dimC = dimA+dimB

C = np.zeros((dimC,dimC,dimC))
for x1 in range(dimA):
    for x2 in range(dimB):
        for y1 in range(dimA):
            for y2 in range(dimB):
                for z1 in range(dimA):
                    for z2 in range(dimB):
                        x = x1+x2
                        y = y1+y2
                        z = z1+z2
                        C[x,y,z] += A[x1,y1,z1] * B[x2,y2,z2] 

Desafortunadamente, esta versión es muy lenta y no utilizable.

Mi segunda versión fue:

C = scipy.signal.fftconvolve(A,B,mode="full")

Pero esto calcula solo los elementos.max(dimA,dimB)

¿Alguien tiene una mejor idea?