Split data to train and test by condition

Question

Suppose I have a pandas DataFrame containing loan information and I would like to predict the probability a user will not return the money (indicated by the default column in my dataframe). I would like to split the data in train and test sets using sklearn.model_selection.train_test_split.

However, I want to make sure that loans with same customerID won't appear both in test and the train set. How should I do this?

Below a sample of my data:

d = {'loan_date': ['20170101','20170701','20170301','20170415','20170515'],
     'customerID': [111,111,222,333,444],
     'loanID': ['aaa','fff','ccc','ddd','bbb'],
     'loan_duration' : [6,3,12,5,12],
     'gender':['F','F','M','F','M'],
     'loan_amount': [20000,10000,30000,10000,40000],
     'default':[0,1,0,0,1]}

df = pd.DataFrame(data=d)

CustomerID==111 loan records, for example, should appear either in the test or the train set, but not in both.

Why's it a problem if you have the same Customer ID in train and test sets? — Josh Friedlander

Mikhail Stepanov Mikhail Stepanov · Accepted Answer · 2019-01-27T20:48:19

I propose following solution. With Customers with same customerID don't appear in train and test; aslo customers splitted by their activity - i.e. approximately equal part of users with the same number of loans will be placed in train and test.

I extend sample of data for the demostration purposes:

d = {'loan_date': ['20170101','20170701','20170301','20170415','20170515','20170905', '20170814', '20170819', '20170304'],         
     'customerID': [111,111,222,333,444,222,111,444,555],        
     'loanID': ['aaa','fff','ccc','ddd','bbb','eee', 'kkk', 'zzz', 'yyy'],                                                         
     'loan_duration' : [6,3,12,5,12, 3, 17, 4, 6],
     'gender':['F','F','M','F','M','M', 'F', 'M','F'],
     'loan_amount': [20000,10000,30000,10000,40000,20000,30000,30000,40000],
     'default':[0,1,0,0,1,0,1,1,0]}

df = pd.DataFrame(data=d)

Code:

from sklearn.model_selection import train_test_split

def group_customers_by_activity(df):
    value_count = df.customerID.value_counts().reset_index()
    df_by_customer = df.set_index('customerID')
    df_s = [df_by_customer.loc[value_count[value_count.customerID == count]['index']] for count in value_count.customerID.unique()]
    return df_s

- this function splits df by customerID activity (number of entries with the same customerID).
Sample output of this function:

group_customers_by_activity(df)
Out:
[           loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 111         20170101    aaa              6      F        20000        0
 111         20170701    fff              3      F        10000        1
 111         20170814    kkk             17      F        30000        1,
            loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 222         20170301    ccc             12      M        30000        0
 222         20170905    eee              3      M        20000        0
 444         20170515    bbb             12      M        40000        1
 444         20170819    zzz              4      M        30000        1,
            loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 333         20170415    ddd              5      F        10000        0
 555         20170304    yyy              6      F        40000        0]

- groups of users with 1, 2, 3 loan(s) etc..

this function splits a group in a manner that user gets to the train or test either:

def split_group(df_group, train_size=0.8):
    customers = df_group.index.unique()
    train_customers, test_customers = train_test_split(customers, train_size=train_size)
    train_df, test_df = df_group.loc[train_customers], df_group.loc[test_customers]
    return train_df, test_df

split_group(df_s[2])
Out:
(           loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 444         20170515    bbb             12      M        40000        1
 444         20170819    zzz              4      M        30000        1,
            loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 222         20170301    ccc             12      M        30000        0
 222         20170905    eee              3      M        20000        0)

The rest is apply this to all the groups of "customer activity":

def get_sized_splits(df_s, train_size):
    train_splits, test_splits = zip(*[split_group(df_group, train_size) for df_group in df_s])
    return train_splits, test_splits

df_s = group_customers_by_activity(df)
train_splits, test_splits = get_sized_splits(df_s, 0.8)
train_splits, test_splits
Out:
((Empty DataFrame
  Columns: [loan_date, loanID, loan_duration, gender, loan_amount, default]
  Index: [],
             loan_date loanID  loan_duration gender  loan_amount  default
  customerID                                                             
  444         20170515    bbb             12      M        40000        1
  444         20170819    zzz              4      M        30000        1,
             loan_date loanID  loan_duration gender  loan_amount  default
  customerID                                                             
  333         20170415    ddd              5      F        10000        0),
 (           loan_date loanID  loan_duration gender  loan_amount  default
  customerID                                                             
  111         20170101    aaa              6      F        20000        0
  111         20170701    fff              3      F        10000        1
  111         20170814    kkk             17      F        30000        1,
             loan_date loanID  loan_duration gender  loan_amount  default
  customerID                                                             
  222         20170301    ccc             12      M        30000        0
  222         20170905    eee              3      M        20000        0,
             loan_date loanID  loan_duration gender  loan_amount  default
  customerID                                                             
  555         20170304    yyy              6      F        40000        0))

Don't fear the emty DataFrame, it will be concatenated soon. The split function has the following definition:

def split(df, train_size):
    df_s = group_customers_by_activity(df)
    train_splits, test_splits = get_sized_splits(df_s, train_size=train_size)
    return pd.concat(train_splits), pd.concat(test_splits)

split(df, 0.8)
Out[106]: 
(           loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 444         20170515    bbb             12      M        40000        1
 444         20170819    zzz              4      M        30000        1
 555         20170304    yyy              6      F        40000        0,
            loan_date loanID  loan_duration gender  loan_amount  default
 customerID                                                             
 111         20170101    aaa              6      F        20000        0
 111         20170701    fff              3      F        10000        1
 111         20170814    kkk             17      F        30000        1
 222         20170301    ccc             12      M        30000        0
 222         20170905    eee              3      M        20000        0
 333         20170415    ddd              5      F        10000        0)

- so, a customerID is placed either in train or test data. I guess such a strang slit (train > test) because of small size of input data.
If you need no grouping by "customerID activity", you can omit it and just use split_group to hit the goal.

Split data to train and test by condition

1 Answers