Julia中的多维差异/梯度

Question

我正在寻找一种在 Julia 中计算多维数组导数的有效方法。准确地说，我想在 Julia 中有一个相当于 numpy.gradient 的东西。但是，Julia 函数 diff ：

• 仅适用于二维数组
• 沿微分维度将数组大小减一

扩展 Julia 的 diff 的定义很简单，因此它可以在 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 维实现，它使用了参数类型和 splatting 等 Julia 特性：

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.

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

Answer 2

更简单的方法：

mydiff(A::AbstractArray,dim) = mapslices(diff, A, dim)

但不确定这在速度方面会如何比较。

编辑：可能会稍微慢一些，但这是将函数扩展到高阶数组的更通用解决方案：

julia> using BenchmarkTools

julia> 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
mydiff (generic function with 1 method)

julia> X = randn(500,500,500);

julia> @benchmark mydiff($X,3)
BenchmarkTools.Trial: 
  samples:          3
  evals/sample:     1
  time tolerance:   5.00%
  memory tolerance: 1.00%
  memory estimate:  2.79 gb
  allocs estimate:  22
  minimum time:     2.05 s (15.64% GC)
  median time:      2.15 s (14.62% GC)
  mean time:        2.16 s (11.05% GC)
  maximum time:     2.29 s (3.61% GC)

julia> @benchmark mapslices(diff,$X,3)
BenchmarkTools.Trial: 
  samples:          2
  evals/sample:     1
  time tolerance:   5.00%
  memory tolerance: 1.00%
  memory estimate:  1.99 gb
  allocs estimate:  3750056
  minimum time:     2.52 s (7.90% GC)
  median time:      2.61 s (9.17% GC)
  mean time:        2.61 s (9.17% GC)
  maximum time:     2.70 s (10.37% GC)