title: How can I split a dataframe using sklearn train test split such that there are equal proportions for each category?
body: <p>I have a dataset with n independent variables and a categorical variable that I would like to perform a regression analysis on. The number of rows of data is different for each category. I would like to split the dataset into test and train data sets such that each category has an equivalent train test split, e.g. 80% to 20%. Here is a simplified reproducible example of what I'm doing.</p>
  <pre><code>import pandas as pd
  import string 
  import numpy as np
  
  from sklearn.model_selection import train_test_split
  
  nrows=1000
  
  cat_values = ['A','B','C','D']
  # defining the category names
  cats = np.random.choice(cat_values,  size=(nrows))
  
  # creating a random dataframe
  df = pd.DataFrame(np.random.randint(0,1000,size=(nrows, 3)), columns=['variable 1','variable 2','variable 3'])
  df['category'] = cats
  
  y = np.random.rand(nrows)
  
  # using sklearn to split into training and test datasets
  X_train, X_test, y_train, y_test = train_test_split(df, y, test_size = .2, random_state =0)
  
  # printing the number of rows in the output training data set for each category 
  for i in range(len(cat_values)):
      print (&quot;number of rows in category &quot; + str(cat_values[i]) + &quot;: &quot; +  str(len(X_train[X_train['category']==cat_values[i]])))
  </code></pre>
  <p>Output:</p>
  <pre><code>number of rows in category A: 221
  number of rows in category B: 188
  number of rows in category C: 179
  number of rows in category D: 212
  </code></pre>
  <p>I would like the rows to be split e.g. 80:20 train:test for each categorical variable. I've looked at using StratifiedShuffleSplit (<a href="https://stackoverflow.com/questions/64353575/train-test-split-preserving-class-proportions-in-each-split">Train/test split preserving class proportions in each split</a>) but there doesn't seem to be an option for specifying which column to stratify the split on (<a href="https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedShuffleSplit.html" rel="nofollow noreferrer">https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.StratifiedShuffleSplit.html</a>).</p>
  <p>Is there a package that can split the data this way, or would I have to divide my dataframe into n categorical dataframes and perform a different train test split on each one before rejoining them?</p>
  <p>Thanks for any assistance with this.</p>