Julia 中的多维 diff/gradient

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 中是否有一种我错过的规范方法来做到这一点？如果有 none，什么是最优的？

谢谢。

Answer 1

我不太熟悉 diff，但根据我对它所做的了解，我做了一个 n 维实现，它使用了 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)

Julia 中的多维 diff/gradient

Multi-dimensional diff/gradient in Julia

numerical-methods

multidimensional-array

julia