multidimensional-array - Julia 中的多维差异/梯度

Question

我正在寻找一种有效的方法来计算 Julia 中多维数组的导数。准确地说，我想要一个等效的 numpy.gradient在 Julia 。但是，Julia 函数 diff :

仅适用于二维数组

沿微分维度将数组的大小减一

扩展 diff 的定义很简单Julia 因此它可以在 3 维数组上工作，例如和

function diff3D(A::Array, dim::Integer)
    if dim == 1
        [A[i+1,j,k] - A[i,j,k] for i=1:size(A,1)-1, j=1:size(A,2), k=1:size(A,3)]
    elseif dim == 2
       [A[i,j+1,k] - A[i,j,k] for i=1:size(A,1), j=1:size(A,2)-1, k=1:size(A,3)]
    elseif dim == 3
       [A[i,j,k+1] - A[i,j,k] for i=1:size(A,1), j=1:size(A,2), k=1:size(A,3)-1]
    else
        throw(ArgumentError("dimension dim must be 1, 2, or 3 got $dim"))
    end
end

这将与例如

a = [i*j*k for i in 1:10, j in 1:10, k in 1:20]

但是，不可能扩展到任意维度，并且不考虑边界，因此梯度可以与原始数组具有相同的维度。

我有一些想法可以在 Julia 中实现 numpy 梯度的类似物，但我担心它们会非常缓慢和难看，因此我的问题是:在 Julia 中是否有一种规范的方法可以做到这一点我错过了？如果没有，什么是最佳的？

谢谢。

Answer 1

我不太熟悉 diff ，但根据我对它所做的事情的理解，我做了一个 n 维实现，它使用了 Julia 的特性，比如参数类型和 splatting:

function mydiff{T,N}(A::Array{T,N}, dim::Int)
    @assert dim <= N
    idxs_1 = [1:size(A,i) for i in 1:N]
    idxs_2 = copy(idxs_1)
    idxs_1[dim] = 1:(size(A,dim)-1)
    idxs_2[dim] = 2:size(A,dim)
    return A[idxs_2...] - A[idxs_1...]
end

进行一些健全性检查:

A = rand(3,3)
@assert diff(A,1) == mydiff(A,1)  # Base diff vs my impl.
@assert diff(A,2) == mydiff(A,2)  # Base diff vs my impl.

A = rand(3,3,3)
@assert diff3D(A,3) == mydiff(A,3)  # Your impl. vs my impl.

请注意，有更多神奇的方法可以做到这一点，例如使用代码生成将专用方法制作到有限维度，但我认为这可能不需要获得足够好的性能。