Последовательный, кумулятивный расчет
Мне нужно сделать подсчет временных рядов, где значение, вычисленное в каждой строке, зависит от результата, вычисленного в предыдущей строке. Я надеюсь использовать удобство 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()
.