Training and test set with respect to group affiliation

Question

I’m using the following function in R to split subjects/samples into training and test set and it works really fine. However, in my dataset the subjects are divided into 2 groups (Patients and control subjects) and therefore, I wish to split the data while maintaining the proportions of patients and control subjects in each training and test set at the same ratio as in the complete data set. How can I do that in R? How can I modify the following function so that it will take into account group affiliation as it split the data into training and test set?

# splitdf function will return a list of training and testing sets#
splitdf <- function(dataframe, seed=NULL) { 
  if (!is.null(seed)) 
     set.seed(seed)

  index <- 1:nrow(dataframe)
  trainindex <- sample(index, trunc(length(index)/2))
  trainset <- dataframe[trainindex, ] 
  testset <- dataframe[-trainindex, ] 
  list(trainset=trainset,testset=testset) 
}

# apply the function
splits <- splitdf(Data, seed=808)

# it returns a list - two data frames called trainset and testset
str(splits)    

# there are "n" observations in each data frame
lapply(splits,nrow)   

# view the first few columns in each data frame
lapply(splits,head)   

# save the training and testing sets as data frames
training <- splits$trainset
testing <- splits$testset`

#

Example: use the built in iris data and split the dataset into training and testing sets. This dataset has 150 samples and has a factor called Species consisting of 3 levels (setosa, versicolor and virginica)

load the iris data

data(iris)

splits the dataset into training and testing sets:

splits <- splitdf(iris, seed=808)

str(splits)
lapply(splits,nrow)
lapply(splits,head)
training <- splits$trainset
testing <- splits$testset

As you can see here, the function “splitdf” does not take into account group affiliation “Species” when it splits the data into training and test set and as the result the number samples with respect to setosa, versicolor and virginica in the training and test set are Not proportional to that of the main dataset. So, How can I modify the function so that it will take into account group affiliation as it split the data into training and test set?

You should give a reproducible example and the expected output. — agstudy
@agstudy Should I upload the head of the dataset to demonstrate how the datset looks like? — Payam Delfani
@agstudy Thanks for clarifying that point. Now I have edited the questions and added one example there to demonstrate the problem. Hope that someone can help me out! — Payam Delfani

Nate Pope Nate Pope · Accepted Answer · 2013-09-22T18:42:55

Here is a solution using plyr with a simulated dataset.

library(plyr)
set.seed(1001)
dat = data.frame(matrix(rnorm(1000), ncol = 10), treatment = sample(c("control", "control", "treatment"), 100, replace = T) )

# divide data set into training and test sets
tr_prop = 0.5    # proportion of full dataset to use for training
training_set = ddply(dat, .(treatment), function(., seed) { set.seed(seed); .[sample(1:nrow(.), trunc(nrow(.) * tr_prop)), ] }, seed = 101)
test_set = ddply(dat, .(treatment), function(., seed) { set.seed(seed); .[-sample(1:nrow(.), trunc(nrow(.) * tr_prop)), ] }, seed = 101)

# check that proportions are equal across datasets
ddply(dat, .(treatment), function(.) nrow(.)/nrow(dat) )
ddply(training_set, .(treatment), function(.) nrow(.)/nrow(training_set) )
ddply(test_set, .(treatment), function(.) nrow(.)/nrow(test_set) )
c(nrow(training_set), nrow(test_set), nrow(dat)) # lengths of sets

Here, I use set.seed() to ensure identical behavior of sample() when constructing the training/test sets with ddply. This strikes me as a bit of a hack; perhaps there is another way to achieve the same result using a single call to **ply (but returning two dataframes). Another option (without egregious use of set.seed) would be to use dlply and then piece together elements of the resulting list into training/test sets:

set.seed(101) # for consistancy with 'ddply' above
split_set = dlply(dat, .(treatment), function(.) { s = sample(1:nrow(.), trunc(nrow(.) * tr_prop)); list(.[s, ], .[-s,]) } )
# join together with ldply()
training_set = ldply(split_set, function(.) .[[1]])
test_set = ldply(split_set, function(.) .[[2]])

Training and test set with respect to group affiliation

load the iris data

splits the dataset into training and testing sets:

1 Answers