Distance matrix along a dimension

Question

Assuming that I want to compute a distance matrix between every element in a vector, I can do that in the following way:

X = np.array([1, 2, 3])
dist = np.triu(np.expand_dims(X, 0) - np.expand_dims(X, 1))
# [[0 1 2]
#  [0 0 1]
#  [0 0 0]]

However, I am not sure how should I do the same if X is a matrix, and I want to compute the pairwise distances for every vector. For example, assuming that I have the following matrix:

X = np.array([[1, 2, 3], [1, 5, 7],[7, 8, 9]])

I would expect to get the following output:

# [[[0 1 2],
#   [0 0 1],
#   [0 0 0]],
#
#  [[0 4 6],
#   [0 0 2],
#   [0 0 0]],
# 
#  [[0 1 2],
#   [0 0 1],
#   [0 0 0]]]

Answer 1

Use np.triu on 3D extended arrays subtracted version -

In [57]: np.triu(X[:,None,:]-X[:,:,None])
Out[57]: 
array([[[0, 1, 2],
        [0, 0, 1],
        [0, 0, 0]],

       [[0, 1, 2],
        [0, 0, 1],
        [0, 0, 0]],

       [[0, 1, 2],
        [0, 0, 1],
        [0, 0, 0]]])

Or to use your trusty np.expand_dims -

np.triu(np.expand_dims(X, 1) - np.expand_dims(X, 2))

Or create a triu mask with something like np.tri and then mask -

mask = ~np.tri(X.shape[-1], dtype=bool)
out = mask*(X[:,None,:]-X[:,:,None])