initial_split() with strata argument behavior

[issactoast]

The problem

I was expecting to have the same dimension of dataframe from
Unexpected size of data frame after initial_split() with strata argument compare to sample_frac()
At least, the original nrow of airquality data is 153 so I guess airquality_frac2 should have 15 or 16.

Reproducible example

#sample_frac and sample_n
library(dplyr)
#> 
#> The following objects are masked from 'package:stats':
#> 
#>     filter, lag
#> The following objects are masked from 'package:base':
#> 
#>     intersect, setdiff, setequal, union

set.seed(2021)
dim(airquality)
#> [1] 153   6

#sample 10% from each month
airquality_frac <- airquality %>% 
    group_by(Month) %>%  
    sample_frac(size = 0.1, replace=FALSE)

library(rsample)

airquality_split <- initial_split(airquality, prop = 0.1,
              strata = "Month")

airquality_frac2 <- training(airquality_split)

dim(airquality_frac)
#> [1] 15  6
dim(airquality_frac2)
#> [1] 18  6

[juliasilge]

This is because of the stratification and how we use floor(n * prop) in strat_sample():

rsample/R/mc.R

Lines 114 to 120 in 54e3235

	strat_sample <- function(x, prop, times, ...) {
	n <- nrow(x)
	idx <- purrr::map(rep(n, times), sample, size = floor(n*prop), ...)
	out <- purrr::map_df(idx, function(ind, x) x[sort(ind), "idx"], x = x)
	out$rs_id <- rep(1:times, each = floor(n*prop))
	out
	}

The quantity airquality$Month is not uniformly distributed in the sample, especially compared to the overall size, so when we do stratified sampling within each month, we get that result you see.

library(rsample)
table(airquality$Month)
#> 
#>  5  6  7  8  9 
#> 31 30 31 31 30

initial_split(airquality, prop = 0.1)
#> <Analysis/Assess/Total>
#> <16/137/153>
initial_split(airquality, prop = 0.9)
#> <Analysis/Assess/Total>
#> <138/15/153>

initial_split(airquality, prop = 0.1, strata = Month)
#> <Analysis/Assess/Total>
#> <18/135/153>
initial_split(airquality, prop = 0.9, strata = Month)
#> <Analysis/Assess/Total>
#> <140/13/153>

^{Created on 2021-02-08 by the reprex package (v1.0.0)}

We get those results like 13 and 18 from adding up the results of multiplying numbers like 31 * 0.9 and then using floor() within the strata.

We don't expect this to have a negative or surprising impact for realistic datasets. Is there a use case where you are finding a problem with this behavior?

[issactoast]

Thank you for the explanation! Yes, it wouldn't affect the serious problem but I think this behavior does not coincide with the user's intention.

prop argument is for the number of samples for analysis based on what I understand. Since the function uses floor() within the strata, the number of samples for analysis should be matched with the floor() function expectation.

library(tidyverse)
library(rsample)

floor(dim(airquality)[1] * 0.1)
#> [1] 15
floor(dim(airquality)[1] * 0.9)
#> [1] 137

floor(table(airquality$Month) * 0.1) %>% sum()
#> [1] 15
floor(table(airquality$Month) * 0.9) %>% sum()
#> [1] 135

So I guess the more natural behavior of initial_split() for with and without strata for Month, should be {15, 137} and {15 and 135}, which also matches with the result from sample_frac() in dplyr package.

[juliasilge]

Your last example is a pretty good example of how this happens:

library(tidyverse)
library(rsample)

initial_split(airquality, prop = 0.1, strata = Month)
#> <Analysis/Assess/Total>
#> <18/135/153>

## this is like the "testing" set
in_size <- floor(table(airquality$Month) * 0.9) %>% sum()
in_size
#> [1] 135

## this is like the "training" set
nrow(airquality) - in_size
#> [1] 18

^{Created on 2021-02-08 by the reprex package (v1.0.0)}

We can see why this becomes a bit surprising at very small sample sizes, especially with stratified sampling, but it is a result of needing to deal with whole numbers for indices within the strata.

[juliasilge]

Let us know if you have further questions on this behavior! 🙌

[issactoast]

Hi @juliasilge, I have been thinking about this issue and stuck for a while. I think I found the code that makes the issue. Could you review the PR and give me your feedback?

[issactoast]

Now here is the result of the previous code from my tweak.

library(rsample)
table(airquality$Month)
#> 
#>  5  6  7  8  9 
#> 31 30 31 31 30

initial_split(airquality, prop = 0.1)
#> <Analysis/Assess/Total>
#> <15/138/153>
initial_split(airquality, prop = 0.9)
#> <Analysis/Assess/Total>
#> <137/16/153>

initial_split(airquality, prop = 0.1, strata = Month)
#> <Analysis/Assess/Total>
#> <15/138/153>
initial_split(airquality, prop = 0.9, strata = Month)
#> <Analysis/Assess/Total>
#> <135/18/153>

[DavisVaughan]

I'm going to try and summarize this problem without the strata argument first.

library(rsample)

n_row <- nrow(airquality)
n_row
#> [1] 153

prop_1 <- 0.1

# Approximate number of rows I expect in my analysis set
# By convention, lets round this down to 15
n_row * prop_1
#> [1] 15.3

# But we get 16?
initial_split(airquality, prop = prop_1)
#> <Analysis/Assess/Total>
#> <16/137/153>

# This is because internally we swap `prop` for `1 - prop` and instead compute
# the assessment set indices first. So the following code corresponds to the
# number of rows in the assessment set. Rounded down this gives us 137
n_row * (1 - prop_1)
#> [1] 137.7

# And this is how we got that 16:
n_row - floor(n_row * (1 - prop_1))
#> [1] 16

Because our documentation explicitly states:

prop | The proportion of data to be retained for modeling/analysis.

I do think that I would consider this to be a bug.

I think that the solution in #219 might be the right way to go. I'll give it a review.

[github-actions]

This issue has been automatically locked. If you believe you have found a related problem, please file a new issue (with a reprex: https://reprex.tidyverse.org) and link to this issue.