2

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])])
Divakar
  • 218,885
  • 19
  • 262
  • 358
xa5i
  • 53
  • 4

1 Answers1

5

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) :

  • We can use 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)).
Divakar
  • 218,885
  • 19
  • 262
  • 358