I am learning numpy, pandas, and sklearn for machine learning now. My teacher gives me code like below, but I guess it is ugly. This code is for splitting the data into training set and testing set, and later we could manipulate the entire data or just training/testing set by the [True, False, True, False ... ] idx_train/idx_test list later.

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

m = 10
x1 = np.arange(m)
x2 = x1 * 2
x3 = x1 * 3
x4 = x1 * 4
x = np.c_[x1, x2, x3, x4]
x = pd.DataFrame(x)

x_train, x_test = train_test_split(x, train_size=0.7, random_state=666)
TRAIN, TEST = 'train', 'test'
pd.options.mode.chained_assignment = None
pd.options.mode.chained_assignment = 'warn'
x = pd.concat([x_train, x_test])
idx_train = x[TRAIN_TEST] == TRAIN
idx_test = np.invert(idx_train)
x_train = x[idx_train]  # later
x_test = x[idx_test]

And I get the below code by myself. I feel it is only a little better.

offset_orig = np.arange(m)
rate = 0.7
offset_train, offset_test = train_test_split(offset_orig, train_size=rate, random_state=666)
if rate < 0.5:
    offset_set_train = set(offset_train)
    idx_train = [i in offset_set_train for i in offset_orig]  # This may be slow. Is numpy broadcast or implicit cycling available here?
    idx_test = np.invert(idx_train)
    offset_set_test = set(offset_test)
    idx_test = [i in offset_set_test for i in offset_orig]  # This may be slow. Is numpy broadcast or implicit cycling available here?
    idx_train = np.invert(idx_test)

x_train = x[idx_train]  # later
x_test = x[idx_test]

I want to get the [True, False, True, False ... ] list idx_train and idx_test for slicing usage later.

I believe that in the world of Python, there should be and will be only one appropriate way to achieve an objective. Could someone give me that very approach to get a "train or test" slicing? (Usage of the lib function train_test_split is required.)

  • \$\begingroup\$ Cross-posted at Stack Overflow \$\endgroup\$
    – Mast
    May 25, 2021 at 9:35
  • \$\begingroup\$ Sorry, I am suggested to post my question here. The original post in StackOverflow has some information entered by other users, and I feel it should not be deleted. \$\endgroup\$
    – cmpltrtok
    May 25, 2021 at 12:38
  • \$\begingroup\$ @Mast The cross post was removed by moderators on SO. Is this still a good question or can it be closed/removed. \$\endgroup\$
    – pacmaninbw
    Aug 3, 2022 at 13:18
  • \$\begingroup\$ @pacmaninbw I see no reason to close it. \$\endgroup\$
    – Mast
    Aug 3, 2022 at 14:07

1 Answer 1


Inspired by Gulzar at Stackoverflow.

I do not realize it works until after a night's sleep.

So, below is enough:

offset_orig = range(m)
offset_train, offset_test = train_test_split(offset_orig, train_size=0.7, random_state=666)

x_train = x.iloc[offset_train]  # later
x_test = x.iloc[offset_test]

The offset_train list ([0, 2, ...]) has the same feature in later programming against the [True, False, True, False ...] list for slicing I pursued at the beginning.

For example, I can use offset_train later as below:

import matplotlib.pyplot as plt
import seaborn as sns
fig, axes = plt.subplots(nrows=2, ncols=2)
fig.set_size_inches(12, 10)
sns.boxplot(data=df.iloc[offset_train], y='count', ax=axes[0][0])
sns.boxplot(data=df.iloc[offset_train], y='count', x='season', ax=axes[0][1])
sns.boxplot(data=df.iloc[offset_train], y='count', x='hour', ax=axes[1][0])
sns.boxplot(data=df.iloc[offset_train], y='count', x='workingday', ax=axes[1][1])

And even as below:

# get noise index
mu_count = df.loc[offset_train, 'count'].mean(axis=0)
sigma_count = df.loc[offset_train, 'count'].std(axis=0)
idx_bad_np = abs(df.loc[offset_train, 'count'] - mu_count) > 3 * sigma_count
idx_good_np = np.invert(idx_bad_np)
# remove noise index
print(f'Before removal: sum(idx good) = {sum(idx_good_np)}, sum(idx bad) = {sum(idx_bad_np)}, len(df[train]) = {len(df.iloc[offset_train])}')
offset_train = list(set(offset_train) - set(np.array(offset_train)[idx_bad_np]))
print(f'After removal: sum(idx good) = {sum(idx_good_np)}, sum(idx bad) = {sum(idx_bad_np)}, len(df[train]) = {len(df.iloc[offset_train])}')
# check again
idx_bad_np = abs(df.loc[offset_train, 'count'] - mu_count) > 3 * sigma_count
print(f'Check sum(idx bad) after removal: {sum(idx_bad_np)}')

Thanks to the community. I feel I have so far to go on the path of Python.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.