Julia:跨观测张量广播成对距离计算

Question

我正在尝试使用 Julia 中的 Distances 包来执行距离矩阵的广播计算。

我了解如何计算某个矩阵X(尺寸为D x N)的单个N x N距离矩阵，其中每个X[:,i] 列存储用于观察i 的D 维特征向量。代码是:

using Distances

dist_matrix = pairwise(Euclidean(), X, dims = 2)

dist_matrix 包含每对 D 维列之间的欧几里德距离，例如dist_matrix[m,n] 存储 X[:,m] 和 X[:,n] 之间的欧氏距离。

现在想象我的数组X实际上是D维观测值的整个张量或“体积”，因此X [:,i,j] 存储我的 D x N 观测值的第 j 个“切片”。因此，整个数组 X 的尺寸为 D x N x T，其中 T 是切片数。

因此，我想计算距离矩阵的张量或“体积”，以便dist_matrix的尺寸为N x N x T。

有没有办法通过在 Julia 中广播 pairwise() 函数来在一行中完成此操作？最快的方法是什么？下面通过基本的 for 循环展示了这个想法:

using Distances

dist_matrix_tensor = zeros(N,N,T);

for t = 1:T
        dist_matrix_tensor[:,:,t] = pairwise(Euclidean(), X[:,:,t], dims = 2)
end

编辑: 我想出了如何使用 mapslices 来做到这一点，但仍然不确定这是否是最好的方法。

using Distances

dist_function(x)  = pairwise(Euclidean(), x, dims = 2) # define a function that gets the N x N distance matrix for a single 'slice'

dist_matrix_tensor = mapslices(dist_function, X, dims = [1,2]) # map your matrix-operating function across the slices of the main tensor X

这当然也可以并行化，因为 X 的每个“切片”在此计算中都是独立的，所以我基本上只是在寻找最快的方法来做到这一点。我总体上也对如何通过广播来做到这一点感兴趣。

Answer 1

如果 X 的维度很大，您使用 mapslices 的解决方案的性能相当不错。下面是 JuliennedArrays 的示例，对于小型 X 来说速度稍快，但当前两个维度大小为 100 时，其性能与 mapslices 相同。

using Distances, JuliennedArrays, BenchmarkTools

dist_function(x)  = pairwise(Euclidean(), x, dims = 2) # define a function that gets the N x N distance matrix for a single 'slice'

X = randn(10,10,20);
dist_matrix_tensor = @btime mapslices(dist_function, X, dims = [1,2]); # 61.172 μs (198 allocations: 42.28 KiB)
dist_matrix_tensor2 = @btime map(dist_function, Slices(X, 1, 2)); # 41.529 μs (62 allocations: 21.67 KiB)

但请注意，JuliennedArrays 返回 Matrix 的 Vector，而不是三维数组。