Home > Software design >  How to compute the kind of distance matrix with vectorization?
How to compute the kind of distance matrix with vectorization?

Time:08-19

I have an numpy array A of shape 4 X 3 X 2. Each line below is a 2D coordinate of a node. (Each three nodes compose a triangle in my finite element analysis.)

array([[[0., 2.],  #node00 
        [2., 2.],  #node01
        [1., 1.]], #node02

       [[0., 2.],  #node10
        [1., 1.],  #node11
        [0., 0.]], #node12

       [[2., 2.],  #node20
        [1., 1.],  #node21
        [2., 0.]], #node22

       [[0., 0.], #node30
        [1., 1.], #node31
        [2., 0.]]]) #node32

I have another numpy array B of coordinates of pre-computed "centers":

array([[1.        , 1.66666667], # center0
       [0.33333333, 1.        ], # center1
       [1.66666667, 1.        ], # center2
       [1.        , 0.33333333]])# center3

How can I efficiently calculate a matrix C of Euclidian distance like this

dist(center0, node00) dist(center0,node01) dist(center0, node02)
dist(center1, node10) dist(center1,node11) dist(center1, node12)
dist(center2, node20) dist(center2,node21) dist(center2, node22)
dist(center3, node30) dist(center3,node31) dist(center3, node32)

where dist represents a Euclidian distance formula like math.dist or numpy.linalg.norm? Namely, the result matrix's i,j element is the distance between center-i to node-ij.

Vectorized code instead of loops is needed, as my actual data is from medical imaging which is very large. With a nested loop, one can obtain the expected output as follows:

In [63]: for i in range(4):
    ...:     for j in range(3):
    ...:         C[i,j]=math.dist(A[i,j], B[i]) 

In [67]: C
Out[67]:
array([[1.05409255, 1.05409255, 0.66666667],
       [1.05409255, 0.66666667, 1.05409255],
       [1.05409255, 0.66666667, 1.05409255],
       [1.05409255, 0.66666667, 1.05409255]])

[Edit] This is different question from Pairwise operations (distance) on two lists in numpy, as things like indexing needs to be properly addressed here.

CodePudding user response:

a = np.reshape(A, [12, 2])
b = B[np.repeat(np.arange(4), 3)]
c = np.reshape(np.linalg.norm(a - b, axis=-1), (4, 3))
c
# array([[1.05409255, 1.05409255, 0.66666667],
#        [1.05409255, 0.66666667, 1.05409255],
#        [1.05409255, 0.66666667, 1.05409255],
#        [1.05409255, 0.66666667, 1.05409255]])
  • Related