拆分训练/测试数据

Question

我需要将数据集划分为训练和验证数据集。 我试着做80-20，但它不符合我的期望。

train_dataset, test_dataset = train_test_split(df, test_size = 0.2, random_state = 0)

我有一个包含 25 个城市的城市变量，每个城市都有几个观察值（在不同的日期）。 我想要的是对每个城市及其观察结果进行 80-20 划分。 我不知道这个方法是否有名字，但我不知道如何 go 关于它。

谢谢你。

Answer 1

这样的事情怎么样？

import numpy as np
from pandas import DataFrame

def split_data(df: DataFrame, ratio: int): 
    length  = len(df)
    indices = list(range(length))
    np.random.shuffle(indices)

    train_indices = indices[:int(length * ratio)]
    test_indices  = indices[round(length * ratio):]

    train_set = df.iloc[train_indices]
    test_set  = df.iloc[test_indices]
    
    return (train_set, test_set)


train, test = split_data(df, 0.80)