编写一个函数将数据分成训练和测试 - 奇怪的错误

Question

夜晚是一个不起眼的帖子，但我似乎无法弄清楚这个功能有什么问题。本质上，想法是编写一个可以输出两个新数据帧的单个函数（一个带有测试，一个带有训练数据），但是我遇到的问题是该函数一次只能输出一个数据帧，所以我试图使用IF命令。

这就是我提出的。

split <- function(df, tr.split, option) { 
set.seed(1337)
ind <- sample(2, nrow(df), replace = TRUE, prob = c(tr.split, 1-tr.split))
if(option=="TEST") {te.data <- df[ind==2,]}
if(option=="TEST") {te.data}
if(option=="TRAIN") {tr.data <- df[ind==1,]}
if(option=="TRAIN") {tr.data}}

相反，当使用我的数据集运行函数时，它返回一个空数据集。它是这样运行的。

te <- split(cleaned.data, 0.8, "TEST")
tr <- split(cleaned.data, 0.8, "TRAIN")

任何帮助将不胜感激 ：）

Answer 1

这个函数提供了两个元素的列表，第一个是训练集，第二个是测试集：

split <- function(df, tr.split) { 
  set.seed(1337)
  ind <- sample(2, nrow(df), replace = TRUE, prob = c(tr.split, 1-tr.split))
  te.data <- df[ind==2,]
  tr.data <- df[ind==1,]

  return(list(tr.data,te.data))
  }

一个例子：

df<-data.frame(x=c(1:20)+round(runif(20,min = 5, max=20),0),
+                date=seq(as.Date("2000/1/1"), by = "day", length.out = 20))
> split(df,0.8)
[[1]]
    x       date
1  19 2000-01-01
2  18 2000-01-02
3   9 2000-01-03
4  11 2000-01-04
5  18 2000-01-05
6  26 2000-01-06
8  26 2000-01-08
9  21 2000-01-09
10 30 2000-01-10
14 22 2000-01-14
16 31 2000-01-16
19 28 2000-01-19
20 28 2000-01-20

[[2]]
    x       date
7  26 2000-01-07
11 29 2000-01-11
12 19 2000-01-12
13 29 2000-01-13
15 32 2000-01-15
17 37 2000-01-17
18 34 2000-01-18