Positive cumulative sum of difference I have a water reservoir with input and output rates. I want to determine when the input is exceeding the output by a certain constant. To accomplish this, I need to cumulatively sum all cases where the inflow exceeds the outflow. Thus, I've written this function:

def pos_diff_cum_sum(flow_in: np.ndarray, flow_out: np.ndarray) -&gt; np.ndarray:
    sums = []
    cum_sum = 0
    diff = list(flow_in - flow_out)

    for dd in diff:
        cum_sum += dd
        if cum_sum &lt; 0:
            cum_sum = 0

        sums.append(cum_sum)

    return np.array(sums)


It sums up the periods where the inflow exceeds the outflow while ignoring periods where the opposite is true. Basically, numpy.cumsum with a corner case.

Tests with plots

t_steps = 9

fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(4, 8))
in_flow = np.linspace(1., 0., t_steps)
out_flow = np.linspace(0., 1., t_steps)

ax1.plot(in_flow, label="in")
ax1.plot(out_flow, label="out")
ax1.legend()

pos_diff = pos_diff_cum_sum(in_flow, out_flow)
ax2.plot(pos_diff)
pos_diff
# =&gt; array([1.  , 1.75, 2.25, 2.5 , 2.5 , 2.25, 1.75, 1.  , 0.  ])




t_steps = 9

fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(4, 8))
in_flow = np.linspace(0., 1., t_steps)
out_flow = np.linspace(1., 0., t_steps)

ax1.plot(in_flow, label="in")
ax1.plot(out_flow, label="out")
ax1.legend()

pos_diff = pos_diff_cum_sum(in_flow, out_flow)
ax2.plot(pos_diff)
pos_diff

# =&gt; array([0.  , 0.  , 0.  , 0.  , 0.  , 0.25, 0.75, 1.5 , 2.5 ])




This code isn't vectorized, but it's going to be called very frequently, so is there some way I should speed it up? Is there any way to make this more elegant?
You can use np.cumsum and np.minimum.accumulate (which I found from this post).

Another way to look at what you want is you want the cumsum.

If the value goes below zero then you want to subtract a value to get it to zero, this value is itself. This means that you just need a running minimum. This is as you've subtracted the value from an earlier value so the effect propagates in your version but not in np.cumsum.

You also want to start this minumum from 0.

def pos_diff_cum_sum(flow_in, flow_out):
    delta = np.cumsum(flow_in - flow_out)
    return delta - np.minimum.accumulate(np.append([0], delta))[1:]




For reference below.

def fn(in_, out):
    delta = np.cumsum(np.array(in_) - np.array(out))
    print(delta)
    output = delta - np.minimum.accumulate(np.append([0], delta))[1:]
    print(np.minimum.accumulate(np.append([0], delta))[1:])
    print(output)


If you have an input that only increases then you can just use use np.cumsum:

&gt;&gt;&gt; fn([1, 1, 1, 1, 1], [0, 0, 0, 0, 0])
[1 2 3 4 5]
[0 0 0 0 0]
[1 2 3 4 5]


However if the number goes negative you must subtract all values after it goes negative by that value. This is as the single -= affects the rest of the input in the OP:

&gt;&gt;&gt; fn([1, 0, 0, 0, 0], [0, 1, 1, 0, 0])
[ 1  0 -1 -1 -1]
[ 0  0 -1 -1 -1]
[1 0 0 0 0]


This means you must subtract them even if the value becomes positive again:

&gt;&gt;&gt; fn([1, 0, 0, 1, 1], [0, 1, 1, 0, 0])
[ 1  0 -1  0  1]
[ 0  0 -1 -1 -1]
[1 0 0 1 2]


If more numbers go negative then you have to decrease by these amounts too:

&gt;&gt;&gt; fn([1, 0, 0, 0, 0], [0, 1, 1, 1, 1])
[ 1  0 -1 -2 -3]
[ 0  0 -1 -2 -3]
[1 0 0 0 0]


This allows the value to go positive again if it needs to:

&gt;&gt;&gt; fn([1, 0, 0, 0, 1], [0, 1, 1, 1, 0])
[ 1  0 -1 -2 -1]
[ 0  0 -1 -2 -2]
[1 0 0 0 1]