15

I'm trying to partition a data set that I have in R, 2/3 for training and 1/3 for testing. I have one classification variable, and seven numerical variables. Each observation is classified as either A, B, C, or D.

For simplicity's sake, let's say that the classification variable, cl, is A for the first 100 observations, B for observations 101 to 200, C till 300, and D till 400. I'm trying to get a partition that has 2/3 of the observations for each of A, B, C, and D (as opposed to simply getting 2/3 of the observations for the entire data set since it will likely not have equal amounts of each classification).

When I try to sample from a subset of the data, such as sample(subset(data, cl=='A')), the columns are reordered instead of the rows.

To summarize, my goal is to have 67 random observations from each of A, B, C, and D as my training data, and store the remaining 33 observations for each of A, B, C, and D as testing data. I have found a very similar question to mine, but it did not factor in multiple variables.

ROMANIA_engineer
  • 54,432
  • 29
  • 203
  • 199
Danny
  • 625
  • 2
  • 8
  • 11

4 Answers4

17

There is actually a nice package caret for dealing with machine learning problems and it contains a function createDataPartition() that pretty much does this sampling 2/3rds from each level of a supplied factor:

#2/3rds for training
library(caret)
inTrain = createDataPartition(df$yourFactor, p = 2/3, list = FALSE)
dfTrain=df[inTrain,]
dfTest=df[-inTrain,]
Stephen Henderson
  • 6,340
  • 3
  • 27
  • 33
5

this may be longer but i think it's more intuitive and can be done in base R ;)

# create the data frame you've described
x <-
    data.frame(
        cl = 
            c( 
                rep( 'A' , 100 ) ,
                rep( 'B' , 100 ) ,
                rep( 'C' , 100 ) ,
                rep( 'D' , 100 ) 
            ) ,

        othernum1 = rnorm( 400 ) ,
        othernum2 = rnorm( 400 ) ,
        othernum3 = rnorm( 400 ) ,
        othernum4 = rnorm( 400 ) ,
        othernum5 = rnorm( 400 ) ,
        othernum6 = rnorm( 400 ) ,
        othernum7 = rnorm( 400 ) 
    )

# sample 67 training rows within classification groups
training.rows <-
    tapply( 
        # numeric vector containing the numbers
        # 1 to nrow( x )
        1:nrow( x ) , 

        # break the sample function out by
        # the classification variable
        x$cl , 

        # use the sample function within
        # each classification variable group
        sample , 

        # send the size = 67 parameter
        # through to the sample() function
        size = 67 
    )

# convert your list back to a numeric vector
tr <- unlist( training.rows )

# split your original data frame into two:

# all the records sampled as training rows
training.df <- x[ tr , ]

# all other records (NOT sampled as training rows)
testing.df <- x[ -tr , ]
Anthony Damico
  • 5,779
  • 7
  • 46
  • 77
  • Brilliant! I've not yet heard of the unlist function. This seems to do exactly what I want, and in a shorter way than what I ended up doing. – Danny Nov 26 '12 at 22:14
4

The following will add a set column with values "train" or "test" to your data.frame:

library(plyr)
df <- ddply(df, "cl", transform, set = sample(c("train", "test"), length(cl),
                                              replace = TRUE, prob = c(2, 1)))

You can get something similar using the base ave function, but I find ddply pretty clean (readable) for this particular usage.

You can then split your data using the subset function:

train.data <- subset(df, set == "train")
test.data  <- subset(df, set == "test")

Follow-up: to split each group exactly into 2/3 and 1/3 sizes, you can use:

df <- ddply(df, "cl", transform,
            set = sample(c(rep("train", round(2/3 * length(cl)),
                           rep("test",  round(1/3 * length(cl)))))
flodel
  • 87,577
  • 21
  • 185
  • 223
  • do you mean transform, not summarize? – frankc Nov 24 '12 at 03:56
  • I may be wrong about this, but it looks like this may not give me exactly 2/3 from each group every time. Isn't this just saying that the probability is 2/3? My problem is that I need 67 from each group every time. – Danny Nov 26 '12 at 22:11
  • @Danny, you are correct. Not a huge deal though, see my edit. – flodel Nov 26 '12 at 22:29
1

Bumped into this issue while constructing my own function for partitioning data for cross-validation with multiple factors for stratification. You could construct such datasets by dividing the data into 3 (or N) equally sized portions while dividing observations within each strata equally to the portions, and then selecting one third as the test set and then combine rest as the training set. I would handle such as list elements in R.

Here is a function that I built using the base package that supports multiple stratification factors, indicated as the column numbers or column names of fields that you wish to have as strata (mtcars dataset example). I think it is rather similar in functionality to ddply, with the exception that you can also use column numbers and that the resulting subsets are given inside a list:

# Function that partitions data into a number of equally (or almost-equally) sized bins that do not overlap, and returns the data bins as a list
# Useful for cross validation
partition_data <- function(
    # Data frame to partition (default example: mtcars data, assuming rows correspond to observations)
    dat = mtcars,
    # Number of equally sized bins to partition to (default here: 2 bins)
    bins = 2,
    # Stratification element, homogeneous subpopulations according to a column that should be subsampled,
    # Observations within a substrata are divided equally to the partitioned bins
    stratum = NA
){
    # Total number of observations
    nobs <- dim(dat)[1]
    # Allocation vector, to be used for randomly distributing the samples to the bins
    loc <- rep(1:bins, times=ceiling(nobs/bins))[1:nobs]


    # If the dataset is stratified, each subpopulation is distributed equally to the bins, otherwise the whole population is the "subpopulation"
    if(missing(stratum)){
        pops <- list(sample(1:dim(dat)[1]))
    }else{
        uniqs <- na.omit(as.matrix(unique(dat[,stratum])))
        pops <- list()
        for(i in 1:nrow(uniqs)){
            # If some of the stratified fields include NA-values, these will not be included in the sampling
            w <- apply(as.matrix(dat[,stratum]), MARGIN=1, FUN=function(x) all(x==uniqs[i,]))
            pops[[i]] <- sample(which(w))
        }
    }
    indices <- vector(length=nobs)
    # Assign the group indices according to permutated samples within each subpopulation
    indices[unlist(pops)] <- loc
    # Assign observations to separate locations in a list
    partitioned_data <- lapply(unique(indices), FUN=function(x) dat[x==indices,])
    # Return the result
    partitioned_data
}

Example of how it works; in this hypothetical example one would wish for the factors 'vs' and 'am' to be equally represented in all the bins:

set.seed(1)

# Stratified sampling, so that combinations of binary covariates vs = {0,1} & am = {0,1} appear equally over the randomized bins of data
pt <- partition_data(mtcars, stratum=c("vs", "am"), bins=3)

# Instances are distributed equally
lapply(pt, FUN=function(x) table(x[,c("vs","am")]))
#> lapply(pt, FUN=function(x) table(x[,c("vs","am")]))
#[[1]]
#   am
#vs  0 1
#  0 4 2
#  1 3 2
#
#[[2]]
#   am
#vs  0 1
#  0 4 2
#  1 2 3
#
#[[3]]
#   am
#vs  0 1
#  0 4 2
#  1 2 2

# 10 or 11 samples (=rows) per partition of data (data had 11 columns)
lapply(pt, FUN=dim)

# Training set containing 2/3 of the stratified samples
# Constructed by dropping out the first third of samples

train <- do.call("rbind", pt[-1])

# Test set containing the remaining 1/3

test <- pt[[1]]

# 21 samples in training dataset
print(dim(train))
# 11 samples in testing dataset
print(dim(test))



> print(train)
                    mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Mazda RX4 Wag      21.0   6 160.0 110 3.90 2.875 17.02  0  1    4    4
Datsun 710         22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1
Hornet 4 Drive     21.4   6 258.0 110 3.08 3.215 19.44  1  0    3    1
Merc 450SE         16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3
Cadillac Fleetwood 10.4   8 472.0 205 2.93 5.250 17.98  0  0    3    4
Fiat 128           32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1
Toyota Corona      21.5   4 120.1  97 3.70 2.465 20.01  1  0    3    1
Dodge Challenger   15.5   8 318.0 150 2.76 3.520 16.87  0  0    3    2
Camaro Z28         13.3   8 350.0 245 3.73 3.840 15.41  0  0    3    4
Ford Pantera L     15.8   8 351.0 264 4.22 3.170 14.50  0  1    5    4
Volvo 142E         21.4   4 121.0 109 4.11 2.780 18.60  1  1    4    2
Hornet Sportabout  18.7   8 360.0 175 3.15 3.440 17.02  0  0    3    2
Duster 360         14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4
Merc 230           22.8   4 140.8  95 3.92 3.150 22.90  1  0    4    2
Merc 280           19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4
Merc 450SLC        15.2   8 275.8 180 3.07 3.780 18.00  0  0    3    3
Honda Civic        30.4   4  75.7  52 4.93 1.615 18.52  1  1    4    2
Pontiac Firebird   19.2   8 400.0 175 3.08 3.845 17.05  0  0    3    2
Porsche 914-2      26.0   4 120.3  91 4.43 2.140 16.70  0  1    5    2
Lotus Europa       30.4   4  95.1 113 3.77 1.513 16.90  1  1    5    2
Ferrari Dino       19.7   6 145.0 175 3.62 2.770 15.50  0  1    5    6
> print(test)
                     mpg cyl  disp  hp drat    wt  qsec vs am gear carb
Mazda RX4           21.0   6 160.0 110 3.90 2.620 16.46  0  1    4    4
Valiant             18.1   6 225.0 105 2.76 3.460 20.22  1  0    3    1
Merc 240D           24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2
Merc 280C           17.8   6 167.6 123 3.92 3.440 18.90  1  0    4    4
Merc 450SL          17.3   8 275.8 180 3.07 3.730 17.60  0  0    3    3
Lincoln Continental 10.4   8 460.0 215 3.00 5.424 17.82  0  0    3    4
Chrysler Imperial   14.7   8 440.0 230 3.23 5.345 17.42  0  0    3    4
Toyota Corolla      33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1
AMC Javelin         15.2   8 304.0 150 3.15 3.435 17.30  0  0    3    2
Fiat X1-9           27.3   4  79.0  66 4.08 1.935 18.90  1  1    4    1
Maserati Bora       15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8


# Example of sampling without stratification; the binary covariates 'vs' and 'am' are probably not distributed equally over the bins
lapply(pt2 <- partition_data(mtcars, bins=3), FUN=function(x) table(x[,c("vs","am")]))

# Stratified according to a single covariate (cylinders)
lapply(pt3 <- partition_data(mtcars, stratum="cyl", bins=3), FUN=function(x) table(x[,c("cyl")]))

In this particular dataset that was discussed, with the data.frame from Anthony's answer:

xpt <- partition_data(x, stratum="cl", bins=3)
# Same as:
#xpt <- partition_data(x, stratum=1, bins=3)

train_xpt <- do.call("rbind", xpt[-1])
test_xpt <- xpt[[1]]
#> summary(train_xpt[,"cl"])
# A  B  C  D 
#67 66 67 67 
#> summary(test_xpt[,"cl"])
# A  B  C  D 
#33 34 33 33 
Teemu Daniel Laajala
  • 2,316
  • 1
  • 26
  • 37