如何使用tidymodels/tidyverse基于多个分层变量对基因生物数据进行初始划分？

嗨，这个问题我之前处理基因数据的时候也遇到过！确实initial_split()默认只支持单个分层变量，但咱们用tidymodels和tidyverse的工具，完全可以搞定多分层的数据拆分，给你分享两个实用的方法～

方法一：组合多个分层变量为复合特征 #

这是最贴合initial_split()原生用法的方式——把你想要分层的多个变量（比如物种、基因功能类型、表达量分组这类基因数据常见特征）合并成一个新的复合变量，用它来做分层依据。

代码示例 #

library(tidymodels)
library(tidyverse)

# 假设你的基因数据框是gene_data，包含species、gene_type、expr_group三个需要分层的变量
gene_data <- gene_data %>%
  # 把多个分层变量拼接成一个新列，用下划线分隔避免歧义
  mutate(stratum_combined = paste(species, gene_type, expr_group, sep = "_"))

# 固定随机种子保证拆分结果可重复
set.seed(123)
# 用复合变量作为strata参数进行拆分，这里设置训练集比例为70%
gene_split <- initial_split(gene_data, strata = stratum_combined, prop = 0.7)
gene_train <- training(gene_split)
gene_test <- testing(gene_split)

注意事项 #

如果组合后的复合类别太多（比如某个类别只有1-2个样本），分层拆分的效果会打折扣。可以先提前检查：

# 查看每个复合类别的样本数
gene_data %>% count(stratum_combined) %>% arrange(n)

对于样本极少的类别，可以考虑合并到相似类别，或者直接将这些小样本全部划入训练集，避免测试集出现极端不平衡的情况。

方法二：分组后手动拆分 #

如果你不想创建新变量，也可以用group_by()结合slice_sample()手动控制每个分层组的拆分比例，这种方式更灵活。

代码示例 #

set.seed(123)
gene_split_df <- gene_data %>%
  # 按所有需要分层的变量分组
  group_by(species, gene_type, expr_group) %>%
  # 给每个样本标记是否属于训练集，这里设置70%的比例
  mutate(is_train = sample(c(TRUE, FALSE), n(), replace = TRUE, prob = c(0.7, 0.3))) %>%
  ungroup()

# 拆分出训练集和测试集
gene_train <- filter(gene_split_df, is_train)
gene_test <- filter(gene_split_df, !is_train)

优化小技巧 #

如果某个分组的样本数太少（比如只有1个），上面的代码可能会把它全分到训练或测试集，你可以加个判断逻辑：

gene_split_df <- gene_data %>%
  group_by(species, gene_type, expr_group) %>%
  mutate(
    is_train = ifelse(n() < 3, TRUE, 
                      sample(c(TRUE, FALSE), n(), replace = TRUE, prob = c(0.7, 0.3)))
  ) %>%
  ungroup()

这里设置样本数少于3的组全部划入训练集，保证测试集的组分布更合理。

拆分后验证 #

不管用哪种方法，拆分后都建议验证一下训练集和测试集的分层变量分布是否一致，确保拆分效果符合预期：

# 对比训练集和测试集的各组样本数比例
bind_rows(
  gene_train %>% count(species, gene_type, expr_group) %>% mutate(set = "训练集"),
  gene_test %>% count(species, gene_type, expr_group) %>% mutate(set = "测试集")
) %>%
  group_by(species, gene_type, expr_group) %>%
  mutate(占比 = n / sum(n)) %>%
  ggplot(aes(x = set, y = 占比, fill = species)) +
  geom_col(position = "dodge") +
  facet_grid(gene_type ~ expr_group) +
  labs(title = "训练集与测试集分层分布对比")

备注：内容来源于stack exchange，提问作者Martingales