scikit-learnのStratifiedShuffleSplitを使ってテストデータを作る方法

マスターデータからトレインデータとテストデータに分ける時、テストデータの目的変数分布がマスターデータの目的変数分布と同じになるように作成する方法の備忘録。

kaggleのタイタニックデータを使用。
csvファイルにはtrainと書いているがこのデータが全データと過程。

import pandas as pd
import numpy as np
titanic_df = pd.read_csv("train.csv", error_bad_lines=False)
titanic_df.head()

output>

今回は簡単のため、数値データのみのデータとするよう処理。

sample_df = titanic_df.drop(columns = ['Name','Sex','Ticket','Cabin','Embarked'])
sample_df = sample_df.dropna()
sample_df.head()

output>

仮にAgeカラムを目的変数としてその他のカラムを説明変数とする。
ここでやりたいのは、テストデータをここから作る際(split)にAgeカラムの年齢分布がマスターデータと同じようになるデータをランダムに抜き出す事。

そのために以下のように年齢を10年刻みに書き換える。ここでは四捨五入を使った。
グラフにすると以下のようになる。

sample_df['Age_rechunk'] = np.round(sample_df['Age'],decimals = -1)
sample_df['Age_rechunk'].where(sample_df['Age_rechunk'] < 71 , 70,inplace=True )
sample_df['Age_rechunk'].hist(bins = 14)

output>

このデータを元にしてscikit-learnのStratifiedShuffleSplitを使う。
データは元データの30%として、先ほど作ったAge_rechunkカラムのデータ分布とテストデータのAge_rechunkカラムのデータ分布が揃うようにテストデータを作成する。

from sklearn.model_selection import StratifiedShuffleSplit
sample = StratifiedShuffleSplit(n_splits = 1,test_size = 0.3, random_state = 20)
for train_,test_ in sample.split(sample_df, sample_df['Age_rechunk']):
    chunk_train = sample_df.loc[train_]
    chunk_test = sample_df.loc[test_]

作成したテストデータをグラフにすると以下のような感じ。
上で作ったマスターデータの分布と同じような感じになってる。

chunk_test['Age_rechunk'].hist(bins = 14)

output>