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

Question

我是一个 "用所有语言写Fortran "的人，试图学习现代编程实践。我有一个一维函数 ft(lx)=HT(x,f(x),lx)，其中 x和 f(x) 是一维数组，大小为 nx和 lx 是输出阵列的大小 ft. 我想申请 HT 在多维数组上 f(x,y,z).

基本上我想申请 HT 在所有三个维度上，从 f(x,y,z) 定义于 (nx,ny,nz) 维网格，以 ft(lx,ly,lz) 定义于 (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.))

可以适用于任意维度？也就是说，它可以根据输入元组和函数的大小自动调整1,2,3维的计算？

Answer 1

我有一段代码 此处 这与你想要的相当接近。关键的工具是 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 的第一维度。C_k. 这将确保下一次迭代的 "循环 "由 @nexprs 对原始数组的第二维进行操作，以此类推。

正如你所看到的，我的代码避免了沿任意维度形成切片，通过换位，我们只需要沿第一个维度进行切片。这使得编程变得更加简单，而且还能带来一些性能上的好处。例如，计算矩阵-向量积 B * C[i1,:,i3] 对于所有 i1,i3可以通过移动第二维度的数据来轻松高效地完成。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)