Pandas 有效地应用依赖于索引值的函数

Question

I have a pandas dataframe with a multindex d, y, b, r, a and I need to apply a function that makes a subtraction among the elements of the dataframe depending on the index value.

To simplify things I will consider only three indices, d, y, r

Index d can take two values value0 and value1 . When d == value0 the index r can only be "O" . Instead when d == value1 r is an integer between 0 and 999. The rest of the indices are common between value0 and value1 .

The dataframe can be constructed as

import numpy as np
import pandas as pd
d_index = [*["value0", "value0"], *["value1" for _ in range(2000)]]

y_index = [*[0, 1], *[0 for _ in range(1000)], *[1 for _ in range(1000)]]

r_index = [*["O", "O"], *[i for i in range(1000)], *[i for i in range(1000)]]

rng = np.random.default_rng(12345)

results00 = rng.uniform(0, 2, 1000).tolist()
results01 = rng.uniform(0, 2, 1000).tolist()

results20 = rng.uniform(0, 6, 1000).tolist()
results21 = rng.uniform(0, 6, 1000).tolist()


variable0 = [1, 2, *results00, *results01]

variable2 = [
    2.5,
    3.5,
    *results20,
    *results21,
]

df = pd.DataFrame(
    {
        "d": d_index,
        "y": y_index,
        "r": r_index,
        "string0": variable0,
        "string2": variable2,
    }
)

df.set_index(["d", "y", "r"], inplace=True)

I need to compute the difference between some columns, such that all indices are the same except for value0 and value1 . The results can be obtained through:

df.loc[("value1", 0), "dstring0"] = (df.loc[("value0", 0), "string0"]).to_numpy() - (
    df.loc[("value1", 0), "string0"]
).to_numpy()

df.loc[("value1", 1), "dstring0"] = (df.loc[("value0", 1), "string0"]).to_numpy() - (
    df.loc[("value1", 1), "string0"]
).to_numpy()

df.loc[("value1", 0), "dstring2"] = (df.loc[("value0", 0), "string2"]).to_numpy() - (
    df.loc[("value1", 0), "string2"]
).to_numpy()

df.loc[("value1", 1), "dstring2"] = (df.loc[("value0", 1), "string2"]).to_numpy() - (
    df.loc[("value1", 1), "string2"]
).to_numpy()

I can deal with this transformation by looping over the y, b, a indices and by performing the subtractions above, however it would not be efficient given the large number of observations (around 8 million)

How can I deal with the operations efficiently?

Edit: added a sample dataframe and the expected output. I also realised that the previous function did not work. I kept it below as a reference

Wrong function

def deviation(df_slice, df, variables):
    d, y, b, r, a = df_slice.name
    dvars = ["d" + var for var in variables]
    if d == "value1":
        df.loc[(d, y, b, r, a), dvars] = (
            df.loc[("value0", y, b, "O", a), variables].to_numpy()
            - df_slice[variables].to_numpy()
        )

df.apply(deviation, axis=1, variables=['string0','string2'],df=df)

Answer 1

The easy answer is swifter. However, it will only improve your performance by a little bit (2 times):

import swifter

df.swifter.apply(deviation, axis=1, variables=['string0','string2'],df=df)

Check out this article from towardsdatascience.com about better pandas vectorization concepts.