How can a function that operates on 2D arrays (cdist) be applied along an axis of a 3D array ?
I tried with numpy.apply_along_axis , but I need to operate on 2D arrays, not 1D. I obtained the result I need by iterating along one axis, but I would prefer it vectorized if possible:
from scipy import spatial
import numpy as np
a = np.random.randn(600).reshape(10, 20, 3)
distances = np.array([spatial.distance.cdist(a[i,:,:], a[i,:,:]) for i in range(a.shape[0])])
Inspired by this post, we can solve it in a vectorized manner. So, following the wiki contents from eucl_dist package (disclaimer: I am its author), we could leverage matrix-multiplication and some NumPy specific implementations, like so -
a_s = np.einsum('ijk,ijk->ij',a,a)
sq_dists = a_s[...,None]+a_s[...,None,:]-2*np.einsum('ijk,ilk->ijl',a,a)
dists = np.sqrt(sq_dists)
Alternative(s) :
np.matmul/@-operator on Python3.x to replace the matrix-multiplication part. Hence, np.einsum('ijk,ilk->ijl',a,a) could be replaced by np.matmul(a,a.transpose(0,2,1)).