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

@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