R语言：用向量化优化四层循环的矩阵求和计算以提升效率

问题背景 #

当前通过四层for循环计算如下求和式：
$$f_{ij} = \sum_{k=1}^9 \sum_{l=1}^9 h_{ik} \cdot h_{jl} \cdot g_{kl}$$
其中样本空间为1-9的整数，i、j取值范围同样为1-9，f、g、h均为9×9矩阵。h矩阵固定，需多次模拟g矩阵（从Dirichlet分布采样），筛选出使sum((y - f)^2)最小的g。当前1000次模拟耗时约1秒，百万次模拟耗时过长，需要通过向量化操作优化计算效率。

原实现代码 #

sim <- function(y, nreps, h) {

  G <- vector("list", nreps)     # 存储Dirichlet分布采样得到的g矩阵
  F <- vector("list", nreps)     # 存储计算得到的f矩阵
  M <- vector("numeric", nreps)  # 存储每次模拟的误差平方和

  require(gtools)

  for(n in 1:nreps) {
    f <- matrix(0, nrow=9, ncol=9)        # 初始化f矩阵
    g <- gtools::rdirichlet(9, rep(1,9))  # 模拟g矩阵
    
    for(i in 1:9) {
      for(j in 1:9) {
        for(k in 1:9) {
          for(l in 1:9) {
            f[i,j] <- f[i,j] + h[i,k] * h[j,l] * g[k,l]  # 求和计算
          }
        }
      }
    }
    F[[n]] <- f  # 保存f矩阵
    G[[n]] <- g  # 保存g矩阵
    M[n] <- sum((y - f)^2) # 计算y与f的误差平方和
  }
  m <- which.min(M)  # 找到误差最小的模拟项
  return(list(g=G[[m]], m=M[m]))
}

调用方式：

sim(y=f.y1, nreps=1000, h=x)

相关数据：

# f.y1数据结构
structure(c(0.0182002022244692, 0.0121334681496461, 0.0101112234580384, 
0, 0, 0, 0, 0, 0, 0.0485338725985844, 0.0940343781597573, 0.112234580384226, 
0.0434782608695652, 0.00910010111223458, 0.00101112234580384, 
0, 0, 0, 0.0333670374115268, 0.110212335692619, 0.132457027300303, 
0.0808897876643074, 0.0222446916076845, 0.0070778564206269, 0.00101112234580384, 
0, 0, 0.0070778564206269, 0.0202224469160768, 0.0596562184024267, 
0.0616784630940344, 0.0262891809908999, 0.0070778564206269, 0, 
0, 0, 0.00202224469160768, 0.00505561172901921, 0.0151668351870576, 
0.0182002022244692, 0.0111223458038423, 0.00404448938321537, 
0, 0, 0, 0.00202224469160768, 0.00404448938321537, 0.00505561172901921, 
0.00505561172901921, 0.00202224469160768, 0.00202224469160768, 
0, 0, 0, 0, 0.00202224469160768, 0.00202224469160768, 0.00202224469160768, 
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
0, 0), class = "table", dim = c(9L, 9L), dimnames = structure(list(
    c("0", "1", "2", "3", "4", "5", "6", "7", "8"), c("0", "1", 
    "2", "3", "4", "5", "6", "7", "8")), names = c("", "")))

# x（即h矩阵）数据结构
structure(c(0.61, 0.16, 0.03, 0.005, 0, 0, 0, 0, 0, 0.32, 0.61, 
0.16, 0.03, 0.005, 0, 0, 0, 0, 0.06, 0.16, 0.61, 0.16, 0.03, 
0.005, 0, 0, 0, 0.01, 0.06, 0.16, 0.61, 0.16, 0.03, 0.01, 0, 
0, 0, 0.01, 0.03, 0.16, 0.61, 0.16, 0.03, 0.01, 0, 0, 0, 0.01, 
0.03, 0.16, 0.61, 0.16, 0.06, 0.01, 0, 0, 0, 0.005, 0.03, 0.16, 
0.61, 0.16, 0.06, 0, 0, 0, 0, 0.005, 0.03, 0.16, 0.61, 0.32, 
0, 0, 0, 0, 0, 0.005, 0.03, 0.16, 0.61), dim = c(9L, 9L))

需提前加载gtools包：

library(gtools)

向量化优化方案 #

核心优化点是将四层循环的求和运算替换为矩阵乘法，完全匹配原求和式的数学逻辑：
$$f = h \cdot g \cdot h^T$$
其中$\cdot$代表矩阵乘法，$h^T$是h的转置矩阵。

同时可进一步优化内存使用：无需存储所有模拟的f和g矩阵，只需实时跟踪当前误差最小的结果，避免不必要的内存占用。

优化后的代码 #

sim_optimized <- function(y, nreps, h) {
  library(gtools)
  
  # 初始化最小误差和对应的g矩阵
  min_m <- Inf
  best_g <- NULL
  
  h_t <- t(h)  # 提前计算h的转置，避免重复计算
  
  for(n in 1:nreps) {
    g <- gtools::rdirichlet(9, rep(1,9))
    # 矩阵乘法替代四层循环
    f <- h %*% g %*% h_t
    current_m <- sum((y - f)^2)
    
    # 仅更新当前最优结果
    if(current_m < min_m) {
      min_m <- current_m
      best_g <- g
    }
  }
  
  return(list(g=best_g, m=min_m))
}

效率对比 #

• 原代码1000次模拟约耗时1秒，优化后的代码耗时可降低至0.01秒以内（具体数值取决于硬件）
• 百万次模拟的耗时将从原有的约15分钟，缩短至约1分钟以内

验证正确性 #

可通过小样本模拟验证优化后结果与原代码一致：

set.seed(123)
res_old <- sim(y=f.y1, nreps=100, h=x)
set.seed(123)
res_new <- sim_optimized(y=f.y1, nreps=100, h=x)

# 检查误差值是否一致
all.equal(res_old$m, res_new$m)
# 检查最优g矩阵是否一致
all.equal(res_old$g, res_new$g)

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