Последовательный, кумулятивный расчет
Мне нужно сделать подсчет временных рядов, где значение, вычисленное в каждой строке, зависит от результата, вычисленного в предыдущей строке. Я надеюсь использовать удобство data.table. Фактической проблемой является гидрологическая модель - расчет совокупного баланса воды, добавление осадков на каждом временном шаге и вычитание стока и испарения в зависимости от текущего объема воды. В набор данных входят различные бассейны и сценарии (группы). Здесь я буду использовать более простую иллюстрацию проблемы.
Упрощенный пример расчета выглядит так: для каждого временного шага (строки) i:
v[i] <- a[i] + b[i] * v[i-1]
a и b - векторы значений параметров, а v - это вектор результата. Для первой строки (i == 1) начальное значение v принимается за v0 = 0.
Первая попытка
Моя первая мысль заключалась в использовании shift() в data.table. Минимальным примером, включая желаемый результат v.ans, является
library(data.table) # version 1.9.7
DT <- data.table(a = 1:4,
b = 0.1,
v.ans = c(1, 2.1, 3.21, 4.321) )
DT
# a b v.ans
# 1: 1 0.1 1.000
# 2: 2 0.1 2.100
# 3: 3 0.1 3.210
# 4: 4 0.1 4.321
DT[, v := NA] # initialize v
DT[, v := a + b * ifelse(is.na(shift(v)), 0, shift(v))][]
# a b v.ans v
# 1: 1 0.1 1.000 1
# 2: 2 0.1 2.100 2
# 3: 3 0.1 3.210 3
# 4: 4 0.1 4.321 4
Это не работает, потому что shift(v) дает копию исходного столбца v, сдвинутого на 1 строку. Это не зависит от назначения v.
Я также подумал о построении уравнения с помощью cumsum() и cumprod(), но это тоже не сработает.
Подход к грубой силе
Поэтому я прибегаю к циклу for внутри функции для удобства:
vcalc <- function(a, b, v0 = 0) {
v <- rep(NA, length(a)) # initialize v
for (i in 1:length(a)) {
v[i] <- a[i] + b[i] * ifelse(i==1, v0, v[i-1])
}
return(v)
}
Эта кумулятивная функция отлично работает с data.table:
DT[, v := vcalc(a, b, 0)][]
# a b v.ans v
# 1: 1 0.1 1.000 1.000
# 2: 2 0.1 2.100 2.100
# 3: 3 0.1 3.210 3.210
# 4: 4 0.1 4.321 4.321
identical(DT$v, DT$v.ans)
# [1] TRUE
Мой вопрос
Мой вопрос: могу ли я написать этот расчет более кратким и эффективным способом data.table, не используя определение for и/или функции? Использование set() возможно?
Или существует лучший подход?
Изменить: лучший цикл
Решение David Rcpp ниже вдохновило меня на удаление ifelse() из цикла for:
vcalc2 <- function(a, b, v0 = 0) {
v <- rep(NA, length(a))
for (i in 1:length(a)) {
v0 <- v[i] <- a[i] + b[i] * v0
}
return(v)
}
vcalc2() на 60% быстрее, чем vcalc().