How to transform my code for speed (for loop) - r

Question

UPDATE: I reduced my code to the pivotal elements to shorten it
The function_impact_calc is very slow (26 secs for 100000 records dataframe). I think the main reason is the for loop (maybe apply or map will help?). Below I simulate the data, write impact_calc function and record run time.
library(dplyr)
library(data.table)
library(reshape2)
###########################################################
# Start Simulate Data
###########################################################
BuySell <- function(m = 40, s = 4) {
S <- pmax(round(rnorm(10, m, s), 2), 0)
S.sorted <- sort(S)
data.frame(buy = rev(head(S.sorted, 5)), sell = tail(S.sorted, 5))
}
number_sates <- 10000
lst <- replicate(number_sates, BuySell(), simplify = FALSE)
# assemble prices data frame
prices <- as.data.frame(data.table::rbindlist(lst))
prices$state_id <- rep(1:number_sates, each = 5)
prices$level <- rep(1:5, times = number_sates)
prices$quantities <- round(runif(number_sates * 5, 100000, 1000000), 0)
# reshape to long format
prices_long <- reshape2::melt(prices,
id.vars = c("state_id", "quantities", "level"),
value.name = "price"
) %>%
rename("side" = "variable") %>%
setDT()
###########################################################
# End Simulate Data
###########################################################
Here is the function impact_calc which is very slow.
##########################################################
# function to optimize
impact_calc <- function(data, required_quantity) {
# get best buy and sell
best_buy <- data[, ,.SDcols = c("price", "side", "level")][side == "buy" & level == 1][1, "price"][[1]]
best_sell <- data[, ,.SDcols = c("price", "side", "level")][side == "sell" & level == 1][1, "price"][[1]]
# calculate mid
mid <- 0.5 * (best_buy + best_sell)
# buys
remaining_qty <- required_quantity
impact <- 0
data_buy <- data[, ,][side == "buy"]
levels <- data_buy[, ,][side == "buy"][, level]
# i think this for loop is slow!
for (level in levels) {
price_difference <- mid - data_buy$price[level]
if (data_buy$quantities[level] >= remaining_qty) {
impact <- impact + remaining_qty * price_difference
remaining_qty <- 0
break
} else {
impact <- impact + data_buy$quantities[level] * price_difference
remaining_qty <- remaining_qty - data_buy$quantities[level]
}
}
rel_impact <- impact / required_quantity / mid
return_list <- list("relative_impact" = rel_impact)
}
The results with run time:
start_time <- Sys.time()
impact_buys <- prices_long[, impact_calc(.SD, 600000), by = .(state_id)]
end_time <- Sys.time()
end_time - start_time
# for 100000 data frame it takes
#Time difference of 26.54057 secs
Thanks for your help!

Answer 1

OP's suspicion is correct: By replacing the for loop by vector operations we can speed up the calculation by a factor of over 100:
required_quantity <- 600000
setDT(prices)
library(bench)
mark(
orig = prices_long[, impact_calc(.SD, required_quantity), by = .(state_id)],
mod1 = prices_long[, impact_calc2(.SD, required_quantity), by = .(state_id)],
vec_w = prices[, {
mid <- 0.5 * (buy[1L] + sell[1L])
tmp <- cumsum(quantities) - required_quantity
list(relative_impact =
sum(pmin(quantities, pmax(0, quantities - tmp)) * (mid - buy)) /
required_quantity / mid)
}, by = .(state_id)],
min_time = 1.0
)
# A tibble: 3 x 13
expression min median `itr/sec` mem_alloc `gc/sec` n_itr n_gc total_time result memory time gc
<bch:expr> <bch:tm> <bch:tm> <dbl> <bch:byt> <dbl> <int> <dbl> <bch:tm> <list> <list> <lis> <list>
1 orig 28.1s 28.1s 0.0356 2.21GB 1.39 1 39 28.1s <data.ta~ <Rprofme~ <bch~ <tibb~
2 mod1 13.1s 13.1s 0.0762 658.42MB 1.45 1 19 13.12s <data.ta~ <Rprofme~ <bch~ <tibb~
3 vec_w 175.1ms 196.9ms 5.19 440.19KB 2.59 6 3 1.16s <data.ta~ <Rprofme~ <bch~ <tibb~
In addition to the speed up, the vectorized version vec_w allocates remarkably less memory (about 5000 times).
Note that the vectorized version vec_w is using the original prices dataset in wide format. So, there is no need to reshape the data from wide to long format.
The second benchmark case mod1 is a version of impact_calc() where the code outside of the for loop code has been modified to make better use of the data.table syntax. These minor modifications alone account for a speed-up by a factor of 2.
The results are identical which is checked by mark().
Explanation of vec_w
If I understand correctly, the OP considers quantities in the given level order until required_quantity is reached. The last level is considered only partiallly to the extent which is required to meet required_quantity exactly.
In a vectorised version this can be achieved by nested ifelse() as shown in this example:
library(data.table)
r <- 5
dt <- data.table(q = 1:4)
dt[, csq := cumsum(q)]
dt[, tmp := csq - r]
dt[, aq := ifelse(tmp < 0, q, ifelse(q - tmp > 0, q - tmp, 0))][]
q csq tmp aq
1: 1 1 -4 1
2: 2 3 -2 2
3: 3 6 1 2
4: 4 10 5 0
The temporary variable tmp holds the difference between the cumulated sum of quantities q and the required quantity r.
The first ifelse() tests if the cumulated sum of quantities q is below the required quantity r. If so then use q without deduction. If not then use the part of q which is required to fill up the cumulated sum of actual quantities aq1 to meet the required quantity r.
The second ifelse() ensures that the quantity q minus deduction is positive (which is the case for the incomplete level) or zero (for the remaining levels below).
The actual quantities aq = c(1, 2, 2, 0) derived by the previous steps do sum up to the requested quantity r = 5.
Now, the ifelse() constructs can be replaced by pmin() and pmax():
dt[, aq := pmin(q, pmax(q - tmp, 0))]
I have verified in a separate benchmark (not posted here) that the pmin()/pmax() approach is slightly faster than the nested ifelse().
Explanation of mod1
In function impact_calc() some code lines can be modified to make use of the data.table syntax.
Thus,
best_buy <- data[, .SD,.SDcols = c("price", "side", "level")][side == "buy" & level == 1][1, "price"][[1]]
best_sell <- data[, .SD,.SDcols = c("price", "side", "level")][side == "sell" & level == 1][1, "price"][[1]]
become
best_buy <- data[side == "buy" & level == 1, first(price)]
best_sell <- data[side == "sell" & level == 1, first(price)]
and
data_buy <- data[, ,][side == "buy"]
levels <- data_buy[, ,][side == "buy"][, level]
become
data_buy <- data[side == "buy"]
levels <- data[side == "buy", level]
I was quite surprised to learn that these modifications outside of the for loop already gained a substantial speed increase.