在R中使函数内的for循环更快

Question

样本数据

set.seed(123)
df <- data.frame(day = 1:365, Precp = sample(1:30, 365, replace = T), 
  ETo = sample(1:10,  365, replace = T), top.FC = 23, CN = 61, DC = 0.4)

该数据包含一年中的某一天、降雨量和蒸散量以及一些常数，例如 top.FC、CN 和 DC。 对于给定的第 i 天，water.update 函数计算第 i 天的土壤水分

water.update <- function(WAT0, RAIN.i, ETo.i, CN, DC, top.FC){ 

S = 25400/CN - 254;  IA = 0.2*S

if (RAIN.i > IA) { RO = (RAIN.i - 0.2 * S)^2/(RAIN.i + 0.8 * S)
} else { 
RO = 0 
}

if (WAT0 + RAIN.i - RO > top.FC) { 
DR = DC * (WAT0 + RAIN.i - RO - top.FC) 
} else { 
DR = 0 
}    
dWAT = RAIN.i - RO - DR - ETo.i
WAT1 = WAT0 + dWAT
WAT1 <- ifelse(WAT1 < 0, 0, WAT1) 
return(list(WAT1,RO,DR))
} 

函数water.model 将water.update 应用于所有天。它是递归的，即每天土壤水都需要前一天的土壤水。因此water.model 函数中的循环。

water.model <- function(dat){

 top.FC  <- unique(dat$top.FC)    

 # I make a vector to store the results 
 dat$WAT <- -9999.9
 dat$RO <- -9999.9
 dat$DR <- -9999.9

# First day (day 1) has a default value
dat$WAT[1] <- top.FC/2 # assuming top soil water is half the content on day 1   
dat$RO[1] <- NA 
dat$DR[1] <- NA

# Now calculate water content for day 2 onwards  

for(d in 1:(nrow(dat)-1)){

 dat[d + 1,7:9] <- water.update(WAT0 = dat$WAT[d], 
                                 RAIN.i = dat$Precp[d + 1], 
                                 ETo.i = dat$ETo[d + 1], 
                                 CN = unique(dat$CN), 
                                 DC = unique(dat$DC),
                                 top.FC = unique(dat$top.FC))
 }
 return(dat)
}


 ptm <- proc.time()
 result <- water.model(df)
 proc.time() - ptm

    user  system elapsed 
    0.18    0.00    0.17 

在这种情况下，for循环是不可避免的，因为它使用前一天的含水量来确定今天的含水量。

有没有更快的方法来编写上述函数？我正在寻找超速 上这段代码。原因是因为我的实际数据要大得多。

Answer 1

使用Rcpp 和data.table。下面的代码运行，但我得到的结果与您提供的 R 代码略有不同。我怀疑这与我解释您使用哪些索引来滞后/领先各个列的方式有关，但是如果没有这些事物所代表的领域知识，我很难直观地理解正确的逻辑应该是什么。希望这是一个不错的起点！

使用以下内容创建一个名为WaterModel.cpp 的单独文件：

#include <Rcpp.h>
using namespace Rcpp;

// [[Rcpp::export]]

List WaterModel(NumericVector RAIN,
                NumericVector ETo,
                double CN,
                double DC,
                double topFC) {

  double S = 25400/CN - 254;
  double IA = 0.2*S;

  int n = RAIN.length();
  NumericVector WAT(n);
  NumericVector RO(n);
  NumericVector DR(n);

  WAT[0] = topFC/2;

  for (int i = 1; i < n; i++) {

    if (RAIN[i] > IA) {
      RO[i] = pow((RAIN[i-1] - 0.2 * S),2) / (RAIN[i-1] + 0.8 * S);
    } else { 
      RO[i] = 0;
    }

    if (WAT[i-1] + RAIN[i-1] - RO[i-1] > topFC) { 
      DR[i] = DC * (WAT[i-1] + RAIN[i-1] - RO[i-1] - topFC) ;
    } else { 
      DR[i] = 0 ;
    } 

    WAT[i] = WAT[i-1] + RAIN[i-1] - RO[i-1] - DR[i-1] - ETo[i-1];

    if (WAT[i] < 0){
      WAT[i] = 0;
    }

  }
    return Rcpp::List::create(Rcpp::Named("WAT") = WAT,
                              Rcpp::Named("RO") = RO,
                              Rcpp::Named("DR") = DR);

}

然后使用Rcpp::sourceCpp() 获取它。然后，您可以将常量保留在 data.table 之外并将它们存储为单个值，而不是对每一行重复它们。当我们真正需要的是单个双精度值时，这使我们无需在 C++ 函数中分配完整向量，并且应该节省一些时间/内存。

library(data.table)
library(Rcpp)

set.seed(123)
DT <- data.table(day = 1:365,
                 Precp = sample(1:30, 365, replace = T), 
                 ETo = sample(1:10,  365, replace = T))

## Don't make constant columns just to store constants
Const_topFC = 23
Const_CN = 61
Const_DC = 0.4

Rcpp::sourceCpp("WaterModel.cpp")

DT[,c("WAT","RO","DR"):= WaterModel(Precp,ETo,Const_CN,Const_DC,Const_topFC)]

DT
#       day Precp ETo     WAT  RO        DR
#   1:   1     9   8 11.50000  0  0.0000000
#   2:   2    24   2 12.50000  0  0.0000000
#   3:   3    13   1 34.50000  0  5.4000000
#   4:   4    27   5 41.10000  0  9.8000000
#   5:   5    29   5 53.30000  0 18.0400000
# ---                                     
# 361: 361     5   8 30.10327  0  8.6166592
# 362: 362     6   9 18.48661  0  4.8413088
# 363: 363    27   7 10.64530  0  0.5946452
# 364: 364    10   8 30.05066  0  5.8581216
# 365: 365    11   1 26.19254  0  6.8202636