使用假设库创建多索引 pd.DataFrame

Question

我需要创建一个pd.DataFrame带有多重索引的。第一个索引级别是一个简单的范围，从1...n. 第二级是日期时间索引。所有列都包含floats. 这是我的示例n=2。from datetime import dateimport pandas as pdfrom hypothesis import givenfrom hypothesis import strategies as stfrom hypothesis.extra.pandas import columns, data_frames, indexes@given(    df1=data_frames(        columns=columns(            ["asset1", "asset2", "asset3", "cash_asset"],            elements=st.floats(allow_nan=False, allow_infinity=False),        ),        index=indexes(            elements=st.dates(                date.fromisoformat("2000-01-01"), date.fromisoformat("2020-12-31")            ),            min_size=10,            unique=True,        ).map(sorted),    ),    df2=data_frames(        columns=columns(            ["asset1", "asset2", "asset3", "cash_asset"],            elements=st.floats(allow_nan=False, allow_infinity=False),        ),        index=indexes(            elements=st.dates(                date.fromisoformat("2000-01-01"), date.fromisoformat("2020-12-31")            ),            min_size=10,            unique=True,        ).map(sorted),    ),)def test_index_level(df1, df2):    df = pd.concat([df1, df2], keys=["df1", "df2"])    assert df.index.nlevels == 2我想知道如何使用hypothesis库直接创建多重索引？很明显，我无法像我的玩具示例中那样手动定义df1、df2等。另一个限制是level 2所有出现的索引都必须相同level 1。

Answer 1

您可以lists组合数据帧，而不是单独定义每个数据帧。

要使每个第一个索引的第二个索引（例如日期）相同，您可以首先计算索引，然后将其提供给生成的列表。也许有一种更简单的方法，但我使用复合材料做到了：

@composite

def df_lists(draw, elements=indexes(

    elements=st.dates(

        date.fromisoformat("2000-01-01"),

        date.fromisoformat("2020-12-31")

    ),

    min_size=10,

    unique=True,

)):

    index = draw(elements.map(sorted))

    df_list = lists(

        data_frames(

            columns=columns(

                ["asset1", "asset2", "asset3", "cash_asset"],

                elements=st.floats(allow_nan=False, allow_infinity=False),

            ),

            index=just(index),  # have to make a strategy from the drawn index values

        ),

        min_size=1, max_size=5  # assume n = 5

    )

    return draw(df_list)

@given(df_lists())

def test_index_level(df_list):

    df = pd.concat(df_list,

                   keys=["df" + str(i + 1) for i in range(len(df_list))])

    assert df.index.nlevels == 2