разделение данных для обучения и случайного тестирования с использованием sklearn

Question

У меня есть файл данных, подобный этому примеру:

пример:

     X  Y month  day  FFMC    DMC     DC   ISI  RH  wind  rain   area
68   2  2   sep  fri  92.4  117.9  668.0  12.2  33   6.3   0.0   0.00
228  4  6   sep  sun  93.5  149.3  728.6   8.1  26   3.1   0.0  64.10
387  5  5   mar  thu  90.9   18.9   30.6   8.0  48   5.4   0.0   0.00

Я пытаюсь разбить его для случайного обучения и проверки наборов, но также на основе столбцов, а не строк. от 3-го столбца до конца и первые 2 столбца будут включены в наборы поездов и тестов. для этого я попытался использовать:

from sklearn.cross_validation import train_test_split
data = pd.read_csv('mydata.txt', sep="\t")
data_train, data_test = train_test_split(data, test_size=0.3)

, но этот пакет разделял строки, а не столбцы. затем я попытался транспонировать файл и использовать такой же пакет, как этот:

X_train, X_test = train_test_split(data.T, test_size=0.3)

вот ожидаемый результат:

поездный набор:

     X  Y   month   day FFMC    DC  ISI RH  area
68   2  2   sep fri 92.4    668.0   12.2    33  0.00
228  4  6   sep sun 93.5    728.6   8.1 26  64.10
387  5  5   mar thu 90.9    30.6    8.0 48  0.00

тестовый набор:

     X  Y   DMC wind    rain
68   2  2   117.9   6.3 0.0
228  4  6   149.3   3.1 0.0
387  5  5   18.9    5.4 0.0

знаете ли вы, как я могу исправить код, чтобы получить ожидаемый набор поездов и тестов?

Answer 1

Вы можете сделать это вручную, написав свою собственную функцию:

import pandas as pd
import numpy as np

data = pd.read_csv('mydata.txt', sep="\t")

columns = data.columns
keep = ['X','Y']

def split_columns(columns_list, keep_columns, frac=0.2):

    # remove the common columns for the moment
    for col in keep_columns:
        columns.remove(col)

    # shuffle the rest of the column list
    np.random.shuffle(columns)

    # select the right proportion for the train and test set
    cut = max(1, int((1-frac)*len(columns)))

    train = columns[:cut]
    test = columns[cut:]

    # Add common columns to both lists 
    train.extend(keep_columns)
    test.extend(keep_columns)

    return train, test

train_columns, test_columns = split_columns(columns, keep)

# Build your train and test set by selecting the appropriate subset of columns
train = data[train_columns]
test = data[test_columns]