Julia 中的去向量化数组操作

Question

这个问题是one of my previous questions 的扩展，内容是在 Python 中优化数组计算。我给定问题的答案是通过@tomjn as 给出的

tot_list = np.where((X == 1) | (Y == 1) | (Z == 1),np.random.random(),1).sum(axis=1)

现在，我想用 Julia 写这个。上面很容易翻译成

tot_list = ifelse.((X.==1).|(Y.==1).|(Z.==1),rand(Float64),1)

请注意，X、Y 和 Z 是逻辑矩阵。当我在 Julia 中运行我的代码时，上述工作正常。但是，它非常缓慢。我正在尝试对上述内容进行去矢量化，我得到了

tot_list = similar(X)
for j in 1:length(tot_list)
    tot_list[j] = ifelse((X[j]==1)|(Y[j]==1)|(Z[j]==1),rand(Float64),1)
end

但是，生成的代码实际上不会产生与矢量化、低效版本相同的结果。我无法举一个例子，因为差异仅在非常大的输入时才明显，但我很肯定上面编写的 Julia 代码的两个部分之间存在很大差异。我究竟做错了什么？两个tot_list不应该一样吗？

编辑：差异可能是由于我在执行此计算之后尝试进行的计算。在我运行上面的代码之前，我定义了

variable_tot = zeros(Float64,1,8*N_state)

运行上面给出的 tot_list 计算后，我再运行

variable_tot=variable_tot+reshape(tot_list,:,length(tot_list))

这和我的问题有关吗？

Answer 1

怎么样（用ifelse编辑，这对我来说更快）

R = rand(size(X)...)  # generate all the random values at once
C = @. X | Y | Z      # the condition
@. ifelse(C, R, 1.0)  # the output

它生成所有随机值（即使是那些你不会使用的随机值，但它会一举完成，我认为这比循环遍历rand() 更快），分配条件C，然后应用@ 987654325@.

编辑2，添加比较：这个解决方案分配更多，但对我来说更快一点：

julia> @btime ((x) -> x ? rand() : 1.0).($X .| $Y .| $Z) ;
  5.598 ms (2 allocations: 7.63 MiB)

julia> @btime (R = rand(size($X)...) ; C = @. $X | $Y | $Z ; @. ifelse(C, R, 1.0)) ;
  3.739 ms (8 allocations: 15.38 MiB)

Answer 2

也许试试这个代码：

((x) -> x ? rand() : 1.0).(X .| Y .| Z)

请注意，对于 Julia 函数而言，类型稳定对于性能至关重要。由于rand() 正在返回Float64，因此替代1 值也需要为Float64 - 因此使用1.0 而不是1。

性能测试 - 我的机器上 1'000'000 个元素矩阵需要 5 毫秒：

X = rand(Bool, 1000, 1000)
Y = rand(Bool, 1000, 1000)
Z = rand(Bool, 1000, 1000)

julia> using BenchmarkTools

julia> @btime ((x) -> x ? rand() : 1.0).($X .| $Y .| $Z)
  5.615 ms (2 allocations: 7.63 MiB)