How do I create train and test samples from one dataframe with pandas?

python

dataframe

train_test_split

pandas

byNikita Barsukov·Dec 3, 2024

To split your DataFrame into training and testing sets in one go, the train_test_split function from sklearn is your friend! Here's how you can apply it:

from sklearn.model_selection import train_test_split

# Magic happening here! ✨
train, test = train_test_split(df, test_size=0.2)

Just setting test_size=0.2 ensures an 80/20 split. Now you're all set to feed this data into your machine learning models.

Comprehensive guide and best practices

While train_test_split is a great tool for a fast split, let's delve deeper into some nuances that can help maximize its utility for various datasets.

Maintaining class equality (Stratified sampling)

For classification problems, you should strive to maintain the proportion of classes in both training and testing sets. Use stratify option for this:

# The pandas DataFrame is now perfectly balanced, as all things should be. 🤓
train, test = train_test_split(df, test_size=0.2, stratify=df['class'])

Ensuring reproducibility (Random seeding)

To ensure reproducibility of your splits, feed a random_state parameter. Trust me, this comes handy when you are comparing model performance:

# Everyone loves reproducibility. It's like having a favorite ritual, but for data science!
train, test = train_test_split(df, test_size=0.2, random_state=42)

Mixing up the data (Shuffling)

If your dataset is potentially influenced by ordering bias, guarantee a random mix of data with shuffle=True (which is the default):

# Let's add some entropy to our split. Chaos is a ladder! 🪜
train, test = train_test_split(df, test_size=0.2, shuffle=True)

Manual split (numpy way)

If you're not a fan of sklearn or prefer a more manual approach, then numpy's randn function is there for you:

import numpy as np

# I prefer doing things manually. My data, my rules! 💪
msk = np.random.rand(len(df)) < 0.8
train = df[msk]
test = df[~msk]

Oh, and don't forget to set the seed for repeatability (because we love reproducibility):

# Again, who doesn't love their seeds?
np.random.seed(42)

Verifying the split (Checks)

After splitting, it's a good practice to check the split to ensure proper distribution:

#Recap time! 📜
print(f"Train set has: {len(train)} records")
print(f"Test set has: {len(test)} records")

Advanced options and clarity

Separate your features and targets

Most of the time, you'll want to separate your predictors from the target variable:

# Let's organize this chaos!
X_train, X_test, y_train, y_test = train_test_split(df.drop('target', axis=1), df['target'], test_size=0.2)

Good ol' pandas way

If you want a pure pandas solution, sample has your back:

# Going back to my roots! 🐼
test = df.sample(frac=0.2, random_state=42)
train = df.drop(test.index)

Data Representativity

No matter your preferred method, always ensure your train and test sets are representative of the whole data. Use sampling techniques if you're dealing with large datasets to avoid memory hiccups.