如何按科目拆分培训数据和测试数据（一个变量）

Question

我正在努力适应逻辑回归。我想按帐户（一个在拟合中不起作用的变量）划分培训和测试数据。我希望他们被帐户分割，每个帐户可以有很多变量。例如，80%的客户将接受培训，20%的客户将接受测试。

我已经尝试了下面的代码，但是这段代码只给了我80%的训练和20%的随机测试。然后在训练数据中，它会给我一些解释，但在测试数据中，它也会给我一个不同变量的精确解释。那不是我想要的。

X_train，X_test，y_train，y_test=train_test_split（X，y，test_size=0.20，random_state=0）

请告知。非常感谢。

Answer 1

这个怎么样

import numpy as np

def group_train_test_split(X, y, test_size, random_state, stratify):
    X = X.copy()
    X['_target'] = y

    X = X.set_index(stratify)
    index = X.index

    index_values = index.unique().values
    np.random.seed(random_state)
    np.random.shuffle(index_values)

    cut = np.round(index_values.shape[0] * test_size).astype('<i4')

    X_test, X_train = X.loc[index_values[:cut]], X.loc[index_values[cut:]]

    return X_test['_target'], X_train['_target'], X_test.drop('_target', axis=1), X_train.drop('_target', axis=1)

y_test, y_train, X_test, X_train = group_train_test_split(X=X, y=y, test_size=0.2, random_state=41, stratify='account')

这样，它将把20%的帐户放入测试数据，而其余的将在训练数据中。