什么是Julia在所需尺寸上点缀产品的方法

Posted 2021-05-06

我有一个带有两个矢量场u和v的xy网格。我将矢量字段表示为Array{Float64, 3}，尺寸为nx×ny×2。我想将点积u.v作为标量场（Array{Float64,2}，尺寸为nx×ny）。实现这一目标的最佳方法是什么？

拥有像dot(u,v,3)这样的东西是完美的，其中3是点积的维度。

nx, ny = 3, 4

u = Array{Float64,3}(rand(0:1, nx, ny, 2))
#[0.0 1.0 0.0 1.0; 1.0 0.0 1.0 0.0; 1.0 0.0 1.0 0.0]
#[1.0 0.0 1.0 0.0; 1.0 0.0 0.0 1.0; 1.0 0.0 1.0 1.0]

v = Array{Float64,3}(rand(0:1, nx, ny, 2))
#[1.0 1.0 1.0 1.0; 0.0 1.0 0.0 0.0; 0.0 1.0 1.0 0.0]
#[1.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0; 1.0 1.0 1.0 0.0]

[dot(u[i,j,:], v[i,j,:]) for i in 1:nx, j in 1:ny]

3×4 Array{Float64,2}:
 1.0  1.0  0.0  1.0
 0.0  0.0  0.0  0.0
 1.0  0.0  2.0  0.0

答案

squeeze(sum(u .* v, 3), 3)干净简洁，但是如果你需要更快的速度，我的电脑速度大约快20倍：

a = fill(zero(eltype(u)), nx, ny)
@inbounds for k in indices(u, 3)
    for j in indices(u, 2)
        for i in indices(u, 1)
            a[i, j] += u[i, j, k] * v[i, j, k]
        end
    end
end

编辑：为了更容易进行基准测试比较，请尝试以下代码：

function dotsum3(u, v)
    a = fill(zero(eltype(u)), size(u, 1), size(u, 2))
    @inbounds for k in indices(u, 3)
        for j in indices(u, 2)
            for i in indices(u, 1)
                a[i, j] += u[i, j, k] * v[i, j, k]
            end
        end
    end
    return a
end

还要注意，如果使用@inbounds，可能应该明确地检查u和v的大小是否兼容。

另一答案

sum(u.*v,3)（合理地）快速和短暂。

要明确获得矩阵，你可以挤压第三维，就像qazxsw poi一样

更新：当然，这有分配，如果速度是一切，不是最好的答案。在这种情况下，请参阅@ DNF的直接循环实现，它基本上与您获得它一样快。

另一答案

另一个版本是

squeeze(sum(u.*v,3), 3)

这很容易扩展到任意第三维长度。我机器的基准测试：

a = copy(view(u,:,:,1))
a .*= view(v,:,:,1)
a .+= @views u[:,:,2].*v[:,:,2]

时间结果：

julia> using BenchmarkTools

julia> function f1(u,v)
       a = fill(zero(eltype(u)), nx, ny)
       @inbounds for k in indices(u, 3)
           for j in indices(u, 2)
               for i in indices(u, 1)
                   a[i, j] += u[i, j, k] * v[i, j, k]
               end
           end
       end
       return a
       end
f1 (generic function with 1 method)

julia> f2(u,v) = squeeze(sum(u .* v, 3), 3)
f2 (generic function with 1 method)

julia> function f3(u,v)
       a = copy(view(u,:,:,1))
       a .*= view(v,:,:,1)
       a .+= @views u[:,:,2].*v[:,:,2]
       return a
       end
f3 (generic function with 1 method)

julia> nx, ny = 3, 4
(3, 4)

julia> u = Array{Float64,3}(rand(0:1, nx, ny, 2));

julia> v = Array{Float64,3}(rand(0:1, nx, ny, 2));

建议的版本更快（考虑到Julia Arrays的内存排序，这是合乎逻辑的）。