I'm trying to do a somewhat complicated pandas groupby operation. Here's some functional but slow pandas code.
# Construct a toy dataframe
idx1 = ["bar", "baz", "foo"]
idx2 = list(range(100, 104))
idx3 = list(range(3))
num_data = len(idx1) * len(idx2) * len(idx3)
index = pd.MultiIndex.from_product((idx1, idx2, idx3), names=["first", "third", "fourth"])
np.random.seed(0)
x = np.random.randint(low=0, high=2, size=num_data, dtype=bool)
input_df = pd.DataFrame(index=index, data={"x": x}).reset_index()
input_df["second"] = "positive"
input_df["second"][input_df["third"] != 100] = "negative"
input_df["third"][input_df["third"] == 101] = 100
# Complication: Not all groups when grouped by "fourth" will have the same indices. Most indices will be shared by most
# "fourth" groups, but the intersection is not complete.
mask = np.ones(num_data, dtype=bool)
mask[[17, 18]] = False
input_df = input_df[mask]
input_df = input_df.set_index(["first", "second", "third", "fourth"])
input_df looks like this:
x
first second third fourth
bar positive 100 0 False
1 False
2 True
negative 100 0 True
1 False
2 True
102 0 False
1 True
2 False
103 0 True
1 False
2 True
baz positive 100 0 False
1 False
2 False
negative 100 0 False # Notice some missing rows here
1 True
102 1 True
2 True
103 0 True
1 True
2 False
foo positive 100 0 False
1 False
2 True
negative 100 0 True
1 False
2 False
102 0 False
1 True
2 True
103 0 True
1 True
2 True
Dataframe guarantees/properties:
- There will always be exactly one positive "third" group in each "first group".
- There are N (variable) negative "third" groups in each "first group"
What I want to do efficiently:
- For each "first" group:
- Compare all negative "third" groups to the single positive "third" group (see code for what "compare" means).
dfs = []
# For for each "first" group:
for first, first_df in input_df.groupby("first"):
# Separate the positive group and negative groups
positive_mask = first_df.index.get_level_values("second") == "positive"
first_df = first_df.droplevel(["first"])
positive_df = first_df[positive_mask]
negative_dfs = first_df[~positive_mask]
positive_df = positive_df.droplevel(["second", "third"])
# Do some computations w.r.t. each negative "third" group and its corresponding positive group.
for third, negative_df in negative_dfs.groupby("third"):
negative_df = negative_df.droplevel(["second", "third"])
# Compare the positive/negative group based on their "fourth" indices alone
# Note that for indices not in their intersection "False" is assigned.
true_true = negative_df["x"] & positive_df["x"]
true_false = negative_df["x"] & ~positive_df["x"]
false_false = ~negative_df["x"] & ~positive_df["x"]
false_true = ~negative_df["x"] & positive_df["x"]
df = pd.DataFrame({
"true_true": true_true,
"true_false": true_false,
"false_false": false_false,
"false_true": false_true
}).reset_index()
df["first"] = first
df["second"] = "negative"
df["third"] = third
dfs.append(df)
# Output: A big dataframe of the computed values of all the negative "third" groups.
output_df = pd.concat(dfs)
output_df = output_df.set_index(["first", "second", "third", "fourth"], verify_integrity=True).sort_index()
That means that output_df looks like this. Note the rows that are all false where there were missing "fourth" indices in the original dataframe.
true_true true_false false_false false_true
first second third fourth
bar negative 100 0 False True False False
1 False False True False
2 True False False False
102 0 False False True False
1 False True False False
2 False False False True
103 0 False True False False
1 False False True False
2 True False False False
baz negative 100 0 False False True False
1 False True False False
2 False False False False # All false from missing data
102 0 False False False False # All false from missing data
1 False True False False
2 False True False False
103 0 False True False False
1 False True False False
2 False False True False
foo negative 100 0 False True False False
1 False False True False
2 False False False True
102 0 False False True False
1 False True False False
2 True False False False
103 0 False True False False
1 False True False False
2 True False False False
Doing this in a loop is extraordinarily slow :( Not only the looping itself, but profiling shows a lot of time being spent doing the inner loop comparison operations since the indices have to be aligned.
Is there a more efficient way to perform this computation, maybe with not so much looping?
EDIT: Added random seed for deterministic example data & updated input/output data.