Julia中高效计算矩阵行两两相关系数的最优方法咨询

假设我有一个名为mat的矩阵：

julia> mat = rand(1:10, 5, 3)
5×3 Matrix{Int64}:
 2  4   3
 5  3  10
 5  7   5
 9  5   7
 4  9   6

我需要计算mat每两行之间的相关系数（例如cor(mat[1, :], mat[2, :])），最终得到一个相关矩阵。目前已编写两种实现脚本并完成基准测试，但希望进一步提升计算速度，以处理2000×20甚至10000×100这类大型数据集。

这是一种直观方法：先用zeros初始化矩阵，再逐一对每两行计算相关系数并填充。但该方法存在冗余计算（例如cor(mat[1, :], mat[3, :])与cor(mat[3, :], mat[1, :])结果完全相同）：

using Statistics

function calc_corr(matrix::Matrix)
    n::Int64 = size(matrix, 1)
    corr_mat = zeros(Float64, n, n)

    for (idx1, idx2)=Iterators.product(1:n, 1:n)
        @inbounds corr_mat[idx1, idx2] = cor(
            view(matrix, idx1, :),
            view(matrix, idx2, :)
        )
    end

    return corr_mat
end

先计算相关矩阵的上三角部分，再通过对称矩阵生成完整结果，避免冗余计算：

using LinearAlgebra

function calc_corr2(matrix::Matrix)
    n::Int64 = size(matrix, 1)
    corr_mat = ones(Float64, n, n)

    # 查找上三角索引
    upper_triang_idx = findall(==(1), triu(ones(Int8, n, n), 1))

    for (idx1, idx2)=Tuple.(upper_triang_idx)
        @inbounds corr_mat[idx1, idx2] = cor(
            view(matrix, idx1, :),
            view(matrix, idx2, :)
        )
    end

    corr_mat = Symmetric(corr_mat)
    return corr_mat
end

1. 小型矩阵测试 #

使用上述mat进行测试：

using BenchmarkTools

@benchmark calc_corr($mat)
BenchmarkTools.Trial: 10000 samples with 10 evaluations.
 Range (min … max):  1.950 μs …   6.210 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     2.160 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   2.178 μs ± 289.600 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▇▇▂ ▁ ▅█▆▂▁▁▃▆▅▃▁ ▁▁▁▁                                      ▂
  ███████████████████████▆█▇▇█▇▇▆▇▆▇▆▆▄▆▄▄▅▃▁▄▄▄▃▅▃▄▄▄▁▁▁▃▄▁▄ █
  1.95 μs      Histogram: log(frequency) by time      3.62 μs <

 Memory estimate: 256 bytes, allocs estimate: 1.

# ---------------------------------------------------------------------

@benchmark calc_corr2($mat)
BenchmarkTools.Trial: 10000 samples with 10 evaluations.
 Range (min … max):  1.220 μs … 773.080 μs  ┊ GC (min … max): 0.00% … 99.19%
 Time  (median):     1.420 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.698 μs ±   9.921 μs  ┊ GC (mean ± σ):  8.15% ±  1.40%

  █          
  █▇▃▂▃▄▂▂▄▄▃▂▃▃▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▁
  1.22 μs         Histogram: frequency by time        3.98 μs <

 Memory estimate: 976 bytes, allocs estimate: 7.

验证两种方案结果一致性：

julia> calc_corr(mat) == calc_corr2(mat)
true

2. 大型矩阵测试 #

test_mat = rand(1:10, 2_000, 20);

@benchmark calc_corr($test_mat)
BenchmarkTools.Trial: 8 samples with 1 evaluation.
 Range (min … max):  632.258 ms … 680.094 ms  ┊ GC (min … max): 0.33% … 1.30%
 Time  (median):     646.215 ms               ┊ GC (median):    0.16%
 Time  (mean ± σ):   650.096 ms ±  16.089 ms  ┊ GC (mean ± σ):  0.49% ± 0.60%

  ▁ ▁           ▁  █       ▁                 ▁                ▁  
  █▁█▁▁▁▁▁▁▁▁▁▁▁█▁▁█▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
  632 ms           Histogram: frequency by time          680 ms <

 Memory estimate: 30.52 MiB, allocs estimate: 2.

# ---------------------------------------------------------------------

@benchmark calc_corr2($test_mat)
BenchmarkTools.Trial: 14 samples with 1 evaluation.
 Range (min … max):  351.040 ms … 396.431 ms  ┊ GC (min … max): 2.58% … 1.81%
 Time  (median):     357.403 ms               ┊ GC (median):    2.86%
 Time  (mean ± σ):   360.863 ms ±  11.661 ms  ┊ GC (mean ± σ):  2.75% ± 0.80%

    █           
  ▇▁█▇▁▇▇▁▁▁▇▁▁▁▇▇▁▇▁▁▇▁▇▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▇ ▁
  351 ms           Histogram: frequency by time          396 ms <

 Memory estimate: 99.63 MiB, allocs estimate: 14.

当前内存不是主要关注点，但面对10000×100这类超大矩阵时，现有方案速度仍不够理想，寻求能进一步提升计算速度的优化建议。

内容的提问来源于stack exchange，提问作者Shayan