Computing weighted mean in pandas dataframe based on intervals - pandas

Question

I have 2 dataframes, df1, df2.
df1 consists of 3 columns, start, end, id.
df2 consists of 4 columns, start, end, id, quantity.
Note that start < end always for both dataframes.
For df1, end - start for each row is always 15, and the [start, end] pair for each row is nonoverlapping and contiguous for each id, e.g.,
df1:
id start end
1 0 15
1 15 30
1 30 45
2 0 15
2 15 30
2 30 45
I need to create a 4th column, quantity_average, in df1, where the quantity_average for each row is the weighted average of all df2.quantity such that the corresponding id is the same in both and there is full/partial overlap between the start, end pairs in both dataframes.
The weight is defined as (min(df2.end, df1.end) - max(df2.start, df1.start)) / 15, i.e., proportional to the amount of overlap.
I will provide a full example. We will use the df1 above, and use
df2 =
id start end quantity
1 0 1.1 3.5
1 1.1 11.4 5.5
1 11.4 34 2.5
1 34 46 3
2 0 1.5 2.2
2 1.5 20 1.0
2 20 30 4.5
So we have the result for quantity_average to be:
1.1 / 15 * 3.5 + (11.4 - 1.1)/15 * 5.5 + (15 - 11.4) / 15 * 2.5 = 4.63333
(30 - 15) / 15 * 2.5 = 2.5
(34 - 30) / 15 * 2.5 = 0.66666
1.5 / 15 * 2.2 + (15 - 1.5) / 15 * 1.0 = 1.12
(20 - 15) / 15 * 1.0 + (30 - 20) / 15 * 4.5 = 3.33333333333
0
I am wondering if there's a quick way to do this in pandas?

Answer 1

Here's one (not so simple) way to do it. It's fast in the sense that it uses vectorized functions but its time and memory complexities are both O(len(df1) * len(df2)). Depending on the scale of your data sets, the memory requirements may overwhelm your computers' hardware.
The idea is to use numpy broadcasting to compare every row in df1 against every row in df2, searching for pairs that:
Have same id
Have overlapping duration (start - end).
... then perform calculations over them:
# Extract the columns to numpy array
# For the columns of df1, raise each by one dimension to prepare
# for numpy broadcasting
id1, start1, end1 = [col[:, None] for col in df1.to_numpy().T]
id2, start2, end2, quantity2 = df2.to_numpy().T
# Match each row in df1 to each row in df2
# `is_match` is a matrix where if cell (i, j) is True, row i of
# df1 matches row j of df2
is_match = (id1 == id2) & (start1 <= end2) & (start2 <= end1)
# `start` is a matrix where cell (i, j) is the maximum start time
# between row i of df1 and row j of df2
start = np.maximum(
np.tile(start1, len(df2)),
np.tile(start2, (len(df1), 1))
)
# Likewise, `end` is a matrix where cell (i, j) is the minium end
# time between row i of df1 and row j of df2
end = np.minimum(
np.tile(end1, len(df2)),
np.tile(end2, (len(df1), 1))
)
# This assumes that every row in df1 has a duration of 15
df1["quantity"] = (is_match * (end - start) * quantity2).sum(axis=1) / 15
# This allow each row in df1 to have a different duration
df1["quantity"] = (is_match * (end - start) * quantity2).sum(axis=1) / (end1 - start1)[:, 0]