0

I have loop to multiplying slices from two 3-d numpy arrays to produce a 2-d array. I suspect there must be a more efficient way to do this with broadcasting or einsum, but I can't figure it out.

A = np.array([[[ 9.99948603e-01, -9.99948603e-01,  0.00000000e+00,
                -0.00000000e+00,  0.00000000e+00, -0.00000000e+00,
                 0.00000000e+00, -0.00000000e+00],
               [ 9.91155487e-08, -3.34504020e-04,  0.00000000e+00,
                -0.00000000e+00,  9.91155487e-08, -3.34504020e-04,
                 0.00000000e+00, -0.00000000e+00]],

              [[ 8.80751720e-01, -9.97142430e-01,  0.00000000e+00,
                -0.00000000e+00,  0.00000000e+00, -0.00000000e+00,
                 8.80751720e-01, -9.97142430e-01],
               [ 9.99948603e-01, -9.99948603e-01,  1.19196784e-01,
                -2.47166897e-03,  8.80751720e-01, -9.97142430e-01,
                 9.91155487e-08, -3.34504020e-04]]])

B = np.array([[[0., 1.],
               [2., 1.]],

              [[0., 1.],
               [1., 0.]]])

C = np.empty((A.shape[-1], B.shape[-1])

for i in range(C.shape[0]):
    for j in range(C.shape[1]):
        C[i, j] = (A[..., i] * B[..., j]).sum()


fgregg
  • 3,173
  • 30
  • 37

0 Answers0