stratified splitting the data

user_3135239 0
Viewed 27221

I have a large data set and like to fit different logistic regression for each City, one of the column in my data. The following 70/30 split works without considering City group.

indexes <- sample(1:nrow(data), size = 0.7*nrow(data))

train <- data[indexes,]
test <- data[-indexes,]

But this does not guarantee the 70/30 split for each city.

lets say that I have City A and City B, where City A has 100 rows, and City B has 900 rows, totaling 1000 rows. Splitting the data with above code will give me 700 rows for train and 300 for test data, but it does not guarantee that i will have 70 rows for City A, and 630 rows for City B in the train data. How do i do that?

Once i have the training data split-ed to 70/30 fashion for each city,i will run logistic regression for each city ( I know how to do this once i have the train data)

logistic-regression r split
5 Answers

The package splitstackshape has a nice function stratified which can do this as well, but this is a bit better than createDataPartition because it can use multiple columns to stratify at once. It can be used with one column like:

library(splitstackshape)
set.seed(42)  # good idea to set the random seed for reproducibility
stratified(data, c('City'), 0.7)

Or with multiple columns:

stratified(data, c('City', 'column2'), 0.7)
user_45496820

Another possible way, similar to IRTFMs answer (e.g., using only base-r) is to use the following. Note that this answer returns a stratified index, which can be used like the index calculated in the question.

p <- 0.7
strats <- your_data$the_stratify_variable

rr <- split(1:length(strats), strats)
idx <- sort(as.numeric(unlist(sapply(rr, function(x) sample(x, length(x) * p)))))

train <- your_data[idx, ]
test <- your_data[-idx, ]

Example:

p <- 0.7
strats <- mtcars$cyl

rr <- split(1:length(strats), strats)
idx <- sort(as.numeric(unlist(sapply(rr, function(x) sample(x, length(x) * p)))))

train <- mtcars[idx, ]
test <- mtcars[-idx, ]

table(mtcars$cyl) / nrow(mtcars)
#>       4       6       8
#> 0.34375 0.21875 0.43750 

table(train$cyl) / nrow(train)
#>    4    6    8
#> 0.35 0.20 0.45 

table(test$cyl) / nrow(test)
#>         4         6         8 
#> 0.3333333 0.2500000 0.4166667

We see that all datasets all (mtcars), train, and test have roughly the same class distributions!

user_30484530
Related