Skip to content

Instantly share code, notes, and snippets.

@tomaskrehlik
Last active December 15, 2015 12:39
Show Gist options
  • Save tomaskrehlik/5262087 to your computer and use it in GitHub Desktop.
Save tomaskrehlik/5262087 to your computer and use it in GitHub Desktop.
Data.table: apply window function by index variable
# Lag-function lags the given variable by the date_variable
lag_variable <- function(data, variable, lags, date_variable = c("Date")) {
if (lags == 0) {
return(data)
}
if (lags>0) {
name <- "lag"
} else {
name <- "forward"
}
require(data.table)
setkeyv(data, date_variable)
if (lags>0) {
data[,index:=seq(1:.N)]
} else {
data[,index:=rev(seq(1:.N))]
}
setkeyv(data, "index")
lags <- abs(lags)
position <- which(names(data)==variable)
for ( j in 1:lags ) {
lagname <- paste(variable,"_",name,j,sep="")
lag <- paste("data[, ",lagname,":=data[list(index-",j,"), ",variable,", roll=TRUE][[",position,"L]]]", sep = "")
eval(parse( text = lag ))
}
setkeyv(data, date_variable)
data[,index:=NULL]
}
# window_func applies the function to the lagged or forwarded variables created by lag_variable
window_func <- function(data, func.name, variable, direction = "window", steps, date_variable = c("Date"), clean = TRUE) {
require(data.table)
require(stringr)
transform <- match.fun(func.name)
l <- length(names(data))
if (direction == "forward") {
lag_variable(data, variable, -steps, date_variable)
cols <- which((!(is.na(str_match(names(a), paste(variable,"_forward(",paste(1:steps,collapse="|"),")",sep=""))[,1])))*1==1)
} else {
if (direction == "backward") {
lag_variable(data, variable, steps, date_variable)
cols <- which((!(is.na(str_match(names(a), paste(variable,"_lag(",paste(1:steps,collapse="|"),")",sep=""))[,1])))*1==1)
} else {
if (direction == "window") {
lag_variable(data, variable, -steps, date_variable)
lag_variable(data, variable, steps, date_variable)
cols <- which((!(is.na(str_match(names(a), paste(variable,"_lag(",paste(1:steps,collapse="|"),")",sep=""))[,1])))*1==1)
cols <- c(cols,which((!(is.na(str_match(names(a), paste(variable,"_forward(",paste(1:steps,collapse="|"),")",sep=""))[,1])))*1==1))
} else {
stop("The direction must be either backward, forward or window.")
}
}
}
data[,transf := apply(data[,cols, with=FALSE], 1, transform)]
if (clean) {
data[,cols:=NULL,with=FALSE]
}
return(data)
}
# Typical use:
# I have a data.table DT with variables Date, value1, value2
# I want to get cumulative sum of next five days
# window_func(DT, "sum", "value1", direction = "forward", steps = 5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment