Обновление кадра данных панд на основе среза

Question

У меня есть фрейм данных, который я хочу разделить на наборы "train" и "test", используя функцию sklearn.model_selection.train_test_split. Эта функция возвращает два среза исходного DataFrame. Однако мне нужно, чтобы это было в одном DataFrame с записью столбца, которая идентифицирует тип записи. Я мог бы написать функцию, которая делает это вместо этого, но использование функции sklearn удобно и надежно.

Мой текущий подход заключается в следующем:

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

dates = pd.date_range('20130101',periods=10)
df = pd.DataFrame(np.random.randn(10,4),index=dates,columns=list('ABCD')).reset_index()

split = [0.8, 0.2]
split_seed = 123

train_df, test_df = model_selection.train_test_split(df, train_size = split[0], test_size = split[1], random_state=split_seed)

train_df["Dataset"] = "train"
test_df["Dataset"] = "test"

final_df = train_df.append(test_df)

Это прекрасно работает, но выдает предупреждение, так как я изменяю скопированные фрагменты вместо исходного df объекта:

A value is trying to be set on a copy of a slice from a DataFrame.
Try using .loc[row_indexer,col_indexer] = value instead

Это не имеет большого значения, так как после этого оригинальный DataFrame больше не используется. Но мне любопытно, как я мог сделать это по-другому. Я предполагаю, что вместо редактирования train_df и test_df и добавления их снова, я мог бы просто отредактировать df напрямую, но, поскольку я не очень хорошо знаком с тем, как работают .loc и .iloc, я изо всех сил пытаюсь посмотрим, как это будет работать.

Код Psuedo, который иллюстрирует то, что я ищу, будет выглядеть следующим образом:

df["Dataset"] = "train" WHERE index in train_df.index.values
df["Dataset"] = "test" WHERE index in test_df.index.values

Answer 1

Один из способов - использовать np.where для добавления условной серии в логическое условие:

df['Dataset'] = np.where(df.index.isin(train_df.index.values), 'train', 'test')

Это предполагает, конечно, что индексы, не содержащиеся в train_df , должны существовать в test_df.

Или используйте np.select для более адаптируемого решения:

conds = [df.index.isin(train_df.index.values),
         df.index.isin(test.index.values)]

df['Dataset'] = np.select(conds, ['train', 'test'], 'other')

Answer 2

Если вы не хотите copy ваш DataFrame при вызове model_selection.train_test_split(), вы можете использовать loc:

df.loc[train_df.index, 'Dataset'] = 'train'
df.loc[test_df.index, 'Dataset'] = 'test'