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

Question

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

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

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] 存储j-th 的“切片”我的D x N 意见。整个数组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 返回的是 Vector 或 Matrix，而不是一个三维数组。