如何构建具有滚动总数的 Python 函数？

Question

我希望构建一个函数，该函数按天为以下 DataFrames 创建滚动总计，其中日期上每个代码的 In 从日期上每个代码的 Out 中减去，这个小计是减去的来自前几天的总数，但总数必须 >=0（我在下面的 Desired Output 中包含了一个示例。

下面是我的输入示例和我正在使用的函数以及我想要的输出示例。

df1 - 在

s = """        Date    Code    Quantity
0   10/01/2019  A   20
3   11/01/2019  A   2
7   12/01/2019  A   4
11  13/01/2019  A   10
"""

df2 - 输出

s ='''    Date     Code   Quantity
0   11/01/2019  A   5
3   12/01/2019  A   100
4   15/01/2019  A   1
6   16/01/2019  A   2
'''

代码

df3 = df1.merge(df2, how='outer', left_on=['date', 'code'], right_on=['date', 'code']).fillna(0)
df3['qty1'] = df3['qty_x'] - df3['qty_y']
df3['qty'] = 0
def final_adder(x):
    x.qty_x = x.qty_x
    print(x)
    return x
df_final = df3.groupby(['code']).apply(final_adder)
df_final['qty'] = df_final['qty'].clip(lower=0)
df_final.drop(['qty_x', 'qty_y','qty1'], inplace=True, axis=1)


          date code  qty_x  qty_y qty1  qty
0   10/01/2019   A   20.0    0.0  20.0    0
3   11/01/2019   A    2.0    5.0  -3.0    0
7   12/01/2019   A    4.0  100.0 -96.0    0
11  13/01/2019   A   10.0    0.0  10.0    0

期望的输出

s = """        Date    Code    Quantity
0   10/01/2019  A   20
3   11/01/2019  A   17
7   12/01/2019  A   0
11  13/01/2019  A   10
12  14/01/2019  A   10
15  15/01/2019  A   9
16  16/01/2019  A   7
"""

Answer 1

关于 SO 的问题有一个完整的子类型，涉及有限制的累积操作（例如：“当累积和变为负数时重置为零”）。这与具有已知重置点（例如，来自另一列，或存在 NaN 等）的累积操作不同，因为条件涉及累积值本身。

在当前的 pandas 或 numpy 中没有干净的方法来以矢量化的方式做到这一点。

我所知道的最好（最快）方法是this SO answer 涉及numba。稍作修改并适应您的问题：

from numba import njit
@njit
def poscumsum(x):
    total = 0
    result = np.empty(x.shape)
    for i, y in enumerate(x):
        total += y
        if total < 0:
            total = 0
        result[i] = total
    return result

有了这个，你可以做到：

a = df1.set_index(['Code', 'Date'])
b = df2.set_index(['Code', 'Date'])
idx = a.index.union(b.index).sort_values()
df3 = (a.reindex(idx, fill_value=0) - b.reindex(idx, fill_value=0))
# optional: resample Date to daily within each group:
df3 = df3.groupby('Code').resample('D', level='Date').sum()
df3['Quantity'] = df3.groupby('Code')['Quantity'].transform(
    lambda g: poscumsum(g.values))

关于问题中提供的数据：

>>> df3
                 Quantity
Code Date                
A    2019-01-10        20
     2019-01-11        17
     2019-01-12         0
     2019-01-13        10
     2019-01-14        10
     2019-01-15         9
     2019-01-16         7

如果您愿意，也可以使用合并。这是一个保留所有中间结果的示例（用于取证分析）：

df3 = df1.merge(df2, on=['Code', 'Date'], how='outer', sort=True).fillna(0)
# optional: resample Date to daily within each group:
df3 = df3.set_index(['Code', 'Date']).groupby('Code').resample('D', level='Date').sum()
df3['diff'] = df3['Quantity_x'] - df3['Quantity_y']
df3['cumdiff'] = df3.groupby('Code')['diff'].transform(
   lambda g: poscumsum(g.values))

df3
# out:
                 Quantity_x  Quantity_y  diff  cumdiff
Code Date                                             
A    2019-01-10        20.0         0.0  20.0     20.0
     2019-01-11         2.0         5.0  -3.0     17.0
     2019-01-12         4.0       100.0 -96.0      0.0
     2019-01-13        10.0         0.0  10.0     10.0
     2019-01-14         0.0         0.0   0.0     10.0
     2019-01-15         0.0         1.0  -1.0      9.0
     2019-01-16         0.0         2.0  -2.0      7.0