如何高效计算DataFrame多列间差值？以JHU新冠病例数据为例

问题背景 #

我正在研究约翰·霍普金斯大学（JHU）的美国及全球COVID-19病例数据集。该数据集以县和州作为行元素，连续日期作为列，日期列中的值为对应县的累计COVID-19病例数。

数据集结构 #

# A tibble: 3,342 x 965
      UID iso2  iso3  code3  FIPS Admin2  Province_State Country_Region   Lat
     <dbl> <chr> <chr> <dbl> <dbl> <chr>   <chr>          <chr>          <dbl>
1 84001001 US    USA     840  1001 Autauga Alabama        US              32.5
2 84001003 US    USA     840  1003 Baldwin Alabama        US              30.7
[内容已省略]
# ... with 956 more variables: Long_ <dbl>, Combined_Key <chr>,
#   `1/22/20` <dbl>, `1/23/20` <dbl>, `1/24/20` <dbl>, `1/25/20` <dbl>,
[内容已省略]
#   `2/11/20` <dbl>, `2/12/20` <dbl>, `2/13/20` <dbl>, `2/14/20` <dbl>, ...

需求 #

为计算每日新增病例数，我尝试了base-R、标准tidyverse及函数式解决方案，发现base-R方法速度快1-3个数量级。请问还有哪些快速计算列间差值的方法？

预期输出 #

# A tibble: 3,178,728 x 5
    FIPS name    value cases date
   <dbl> <chr>   <dbl> <dbl> <date>
 1  1001 1/22/20     0    NA 2020-01-22
 2  1001 1/23/20     0     0 2020-01-23
 3  1001 1/24/20     0     0 2020-01-24

已尝试的三种方法 #

方法1：Tidyverse风格的pivot_longer、group_by()和mutate() #

us_cases3 <-
    us_cases_raw %>%
    select(-c(1:4, 6:11)) %>%
    pivot_longer(!FIPS) %>%
    group_by(FIPS) %>%
    mutate(date = mdy(name),
            cases = value - lag(value))  %>%
    ungroup()
# 耗时：24.52673 秒

# 不使用group_by()的情况：耗时1.519198秒（结果显然错误，但能看出耗时主要在分组操作）

方法2：base-R #

us_cases2 <- us_cases_raw[-c(1:4, 6:11)]
ncol <- ncol(us_cases2)
us_cases3 <- cbind(us_cases2[1], 
                   us_cases2[3:(ncol)] - us_cases2[2:(ncol-1)]) %>% 
             pivot_longer(!FIPS) %>% mutate(date = mdy(name))
# 耗时：1.550864 秒

方法3：purrr::map_*或future::future_map_*函数式方案 #

us_cases4 <-
    us_cases_raw %>%
    select(-c(1:4, 6:11)) %>%
    pivot_longer(!FIPS)

us_cases3 <- map_dfr(unique(us_cases4$FIPS), ~ subset(us_cases4, FIPS == .x) %>% mutate(cases = value - lag(value))) %>%
    mutate(date = mdy(name))
# 耗时：3.801137 分钟

# 启用多进程（加载future包并设置plan(multisession)）
us_cases4 <-
    us_cases_raw %>%
    select(-c(1:4, 6:11)) %>%
    pivot_longer(!FIPS)

us_cases3 <- future_map_dfr(unique(us_cases4$FIPS), ~ subset(us_cases4, FIPS == .x) %>% mutate(cases = value - lag(value))) %>%
    mutate(date = mdy(name))
# 耗时：1.580935 分钟

其他高效计算方法 #

方法4：使用data.table #

data.table针对大数据集做了大量底层优化，行操作和格式转换效率远高于常规tidyverse方法：

library(data.table)

# 转换为data.table格式
setDT(us_cases_raw)
# 提取所有日期格式的列名
date_cols <- grep("^\\d+/\\d+/\\d+$", names(us_cases_raw), value = TRUE)

# 按行计算每日新增病例：第一列设为NA，后续列用当前列减前一列
us_cases_dt <- us_cases_raw[, c(.(FIPS = FIPS), lapply(seq_along(date_cols), function(i) {
    if (i == 1) NA_real_ else .SD[[date_cols[i]]] - .SD[[date_cols[i-1]]]
})), .SDcols = date_cols]

# 给新增病例列命名
setnames(us_cases_dt, c("FIPS", paste0("cases_", date_cols)))

# 转换为长格式，分别处理新增病例和累计病例
cases_long <- melt(us_cases_dt, id.vars = "FIPS", variable.name = "name", value.name = "cases")
cases_long[, name := gsub("cases_", "", name)]
value_long <- melt(us_cases_raw[, .(FIPS, ..date_cols)], id.vars = "FIPS", variable.name = "name", value.name = "value")

# 合并数据并转换日期格式
us_cases_final <- merge(cases_long, value_long, by = c("FIPS", "name"))
us_cases_final[, date := mdy(name)]

方法5：基于矩阵的向量化运算 #

利用矩阵的底层优化实现批量差值计算，完全规避循环或分组操作，是效率最高的方案之一：

# 提取FIPS列和日期对应的累计病例矩阵
fips_vec <- us_cases_raw$FIPS
cases_matrix <- as.matrix(us_cases_raw[, grep("^\\d+/\\d+/\\d+$", names(us_cases_raw))])

# 计算每一行的列间差值：第一列补NA，后续列用当前列减前一列
diff_matrix <- cbind(NA_real_, cases_matrix[, -1] - cases_matrix[, -ncol(cases_matrix)])

# 将矩阵转换为长格式数据框
result_df <- data.frame(
    FIPS = rep(fips_vec, ncol(diff_matrix)),
    name = rep(colnames(diff_matrix), each = nrow(diff_matrix)),
    cases = as.vector(diff_matrix),
    value = as.vector(cases_matrix),
    stringsAsFactors = FALSE
)

# 转换数据类型并处理日期
result_df$FIPS <- as.numeric(result_df$FIPS)
result_df$value <- as.numeric(result_df$value)
result_df$cases <- as.numeric(result_df$cases)
result_df$date <- mdy(result_df$name)

方法6：优化后的tidyverse方案 #

通过c_across实现行内向量操作，避免分组后的lag运算，比原始tidyverse方法效率更高：

library(dplyr)
library(tidyr)
library(purrr)
library(lubridate)

us_cases_optimized <- us_cases_raw %>%
    select(FIPS, matches("^\\d+/\\d+/\\d+$")) %>%
    # 按行处理，计算每一行的列间差值
    rowwise() %>%
    mutate(cases = list(c(NA, diff(c_across(matches("^\\d+/\\d+/\\d+$")))))) %>%
    ungroup() %>%
    # 将差值列表展开
    unnest(c(cases)) %>%
    # 转换为长格式
    pivot_longer(cols = matches("^\\d+/\\d+/\\d+$"), names_to = "name", values_to = "value") %>%
    # 转换日期格式
    mutate(date = mdy(name))

内容的提问来源于stack exchange，提问作者Rob Hanssen