title: Create test/train split based on two groups with Pandas Scikit-learn
body: <p>I have a Pandas dataframe: <code>comb</code>
  the number of <code>ENROLLED_Response</code> entries is quite small, so just random sampling of the entire DataFrame may lose too much of the enrolled data.</p>
  
  <p>the solution is to take a 75% sample of all the entries where <code>ENROLLED_Response == True</code>
  and then to take a 70% sample of all the entries where <code>ENROLLED_Response == False</code></p>
  
  <p>So I should end up with a column <code>is_train</code> with <code>true</code>/<code>false</code> on the DataFrame</p>
  
  <p>so I usually use something like:</p>
  
  <pre><code>from sklearn.cross_validation import cross_val_score
  
  #split the dataset for train and test
  comb['is_train'] = np.random.uniform(0, 1, len(comb)) &lt;= .75
  train, test = comb[comb['is_train']==True],  comb[comb['is_train']==False]
  </code></pre>
  
  <p>Which is fine for most cases,but due to the low number of enrolled this approach tends to leave out too many of the 'enrolled' since there are so few. So what I need is something more like:</p>
  
  <pre><code>comb['is_train'] = train_test_split(comb['ENROLLED_Response']==True, Train_size = 0.75)
  comb['is_train']= train_test_split(comb['ENROLLED_Response']==False, Train_size = 0.75)
  </code></pre>
  
  <p>which of course does not work. the concept is: 1st sample the enrolled and mark a random .75 of them  as train, then sample the non enrolled (everything else) and mark .75 of them as train, in the same new column (is_train) so that  then it can easily be used in Scikit_learn like:</p>
  
  <pre><code>train, test = comb[comb['is_train']==True],comb[comb['is_train']==False]
  </code></pre>
  
  <p>Can't figure out how to do it since the np array generated by random is with respect to the length of the whole DataFrame (among other issues...)</p>