Julia：将一维 Julia 函数应用于多维数组

Question

我是 "write Fortran in all languages" 那种尝试学习现代编程实践的人。我有一个一维函数 ft(lx)=HT(x,f(x),lx)，其中 x 和 f(x) 是大小为 nx 的一维数组，lx 是输出数组的大小 ft。我想在多维数组 f(x,y,z) 上应用 HT。

基本上我想在所有三个维度上应用 HT，从 (nx,ny,nz) 维度网格上定义的 f(x,y,z) 到 (lx,ly,lz) 上定义的 ft(lx,ly,lz)维度网格：

ft(lx,y,z)   = HT(x,f(x,y,z)   ,lx)
ft(lx,ly,z)  = HT(y,ft(lx,y,z) ,ly)
ft(lx,ly,lz) = HT(z,ft(lx,ly,z),lz)

在 f95 风格中我倾向于这样写：

FTx=zeros((lx,ny,nz))
for k=1:nz
for j=1:ny
    FTx[:,j,k]=HT(x,f[:,j,k],lx)
end
end

FTxy=zeros((lx,ly,nz))
for k=1:nz
for i=1:lx
    FTxy[i,:,k]=HT(y,FTx[i,:,k],ly)
end
end

FTxyz=zeros((lx,ly,lz))
for j=1:ly
for i=1:lx
    FTxyz[i,j,:]=HT(z,FTxy[i,j,:],lz)
end
end

我知道惯用的 Julia 需要使用类似 mapslices 的东西。我无法从 mapslices 文档中理解如何去做。

所以我的问题是：什么是惯用的 Julia 代码以及适当的类型声明，相当于 Fortran 风格的版本？

后续的子问题是：是否可以编写函数

FT = HTnD((Tuple of x,y,z etc.),f(x,y,z), (Tuple of lx,ly,lz etc.))

适用于任意维度？ IE。它会根据输入元组和函数的大小自动调整 1,2,3 维的计算？

Answer 1

我有一段代码 here which is fairly close to what you want. The key tool is Base.Cartesian.@nexprs 您可以在链接的文档中阅读。

我的代码中最重要的三行是第 30 行到第 32 行。这里是对它们的作用的口头描述。

第 30 行：将 n1 x n2 x ... nN 大小的数组 C_{k-1} 重塑为 n1 x prod(n2,...,nN) 矩阵 tmp_k.
第 31 行：将函数 B[k] 应用于 tmp_k 的每一列。在我的代码中，这里有一些间接寻址，因为我想允许 B[k] 成为矩阵或函数，但基本思想如上所述。这是您要引入 HT 函数的部分。
第 32 行：将 tmp_k 重塑回 N 维数组并循环排列维度，使得 tmp_k 的第二维最终成为 [=] 的第一维22=]。这确保 @nexprs 隐含的 "loop" 的下一次迭代在原始数组的第二个维度上运行，依此类推。

如您所见，我的代码通过置换避免了沿任意维度形成切片，这样我们只需要沿第一个维度进行切片。这使编程变得更加容易，并且还可以带来一些性能优势。例如，通过将 C 的第二维移动到tmp 并使用 gemm 计算 B * tmp。在没有排列的情况下有效地做同样的事情会困难得多。

Answer 2

按照@gTcV 的代码，您的函数将如下所示：

using Base.Cartesian

ht(x,F,d) = mapslices(f -> HT(x, f, d), F, dims = 1)

@generated function HTnD(
    xx::NTuple{N,Any},
    F::AbstractArray{<:Any,N},
    newdims::NTuple{N,Int}
) where {N}
    quote
        F_0 = F
        Base.Cartesian.@nexprs $N k->begin
            tmp_k = reshape(F_{k-1},(size(F_{k-1},1),prod(Base.tail(size(F_{k-1})))))
            tmp_k = ht(xx[k], tmp_k, newdims[k])
            F_k = Array(reshape(permutedims(tmp_k),(Base.tail(size(F_{k-1}))...,size(tmp_k,1))))
            # https://github.com/JuliaLang/julia/issues/30988
        end
        return $(Symbol("F_",N))
    end
end

一个更简单的版本，显示了 mapslices 的用法，如下所示

function simpleHTnD(
    xx::NTuple{N,Any},
    F::AbstractArray{<:Any,N},
    newdims::NTuple{N,Int}
) where {N}
    for k = 1:N
        F = mapslices(f -> HT(xx[k], f, newdims[k]), F, dims = k)
    end
    return F
end

如果你是单线的朋友，你甚至可以使用 foldl;-)

fold_HTnD(xx, F, newdims) = foldl((F, k) -> mapslices(f -> HT(xx[k], f, newdims[k]), F, dims = k), 1:length(xx), init = F)

Julia：将一维 Julia 函数应用于多维数组

Julia: Apply 1 dimensional Julia function to multi-dimensional array

idioms

multidimensional-array

julia