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

Question

我需要创建一个pd.DataFrame带有多重索引的。第一个索引级别是一个简单的范围，从1...n. 第二级是日期时间索引。所有列都包含floats. 这是我的示例n=2。

from datetime import date

import pandas as pd

from hypothesis import given

from hypothesis import strategies as st

from 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组合数据帧，而不是单独定义每个数据帧。要使每个第一个索引的第二个索引（例如日期）相同，您可以首先计算索引，然后将其提供给生成的列表。也许有一种更简单的方法，但我使用复合材料做到了：@compositedef 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