Apply Function Along Axis Over Two Numpy Arrays - Shapes Not Aligned

I'm probably not seeing something obvious here but don't believe np.apply_along_axis or np.apply_over_axes is what I'm looking for. Say I have the following two arrays: arr1 = np.

Solution 1:

For generic shaped 3D and 2D arrays - arr3 and arr4, we can use some np.einsum magic to have a vectorized solution, like so -

dot1 = np.einsum('ijk,ijl->ikl',arr3,arr3)
dot2 = np.einsum('ijk,ij->ik',arr3,arr4)
inv1 = np.linalg.inv(dot1)
tgt_out = np.einsum('ijk,ij->ik',inv1, dot2)

Runtime test

Approaches -

def org_app(arr3, arr4):
    tgt = []
    for i, j in zip(arr3, arr4):
        tgt.append(coefs(i, j))
    return np.array(tgt)

def einsum_app(arr3, arr4):
    dot1 = np.einsum('ijk,ijl->ikl',arr3,arr3)
    dot2 = np.einsum('ijk,ij->ik',arr3,arr4)
    inv1 = np.linalg.inv(dot1)
    return np.einsum('ijk,ij->ik',inv1, dot2)

Timings and verification -

In [215]: arr3 = np.random.rand(50,50,50)
     ...: arr4 = np.random.rand(50,50)
     ...: 

In [216]: np.allclose(org_app(arr3, arr4), einsum_app(arr3, arr4))
Out[216]: True

In [217]: %timeit org_app(arr3, arr4)
100 loops, best of 3: 4.82 ms per loop

In [218]: %timeit einsum_app(arr3, arr4)
100 loops, best of 3: 19.7 ms per loop

Doesn't look like einsum is giving us any benefits here. This is expected because basically einsum is fighting it out against np.dot, which is much better at sum-reduction and even though we are using it in a loop. The only situation/case in which we can give np.dot a fight is when we loop enough and that should make einsum competitive. We are looping for times equal to the length equal of the first axis of the input arrays. Let's increase it and test again -

In [219]: arr3 = np.random.rand(1000,10,10)
     ...: arr4 = np.random.rand(1000,10)
     ...: 

In [220]: %timeit org_app(arr3, arr4)
10 loops, best of 3: 23 ms per loop

In [221]: %timeit einsum_app(arr3, arr4)
100 loops, best of 3: 9.1 ms per loop

einsum definitely winning on this one!

This related post on the fight between np.einsum and np.dot is worth a look.

Also, note that if we need to use the loop based approach, we should look to initialize the output array and then assign the output values from coefs into it rather than appending, as the latter is a slow process.