6
\$\begingroup\$

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) -> np.ndarray:
    sums = []
    cum_sum = 0
    diff = list(flow_in - flow_out)

    for dd in diff:
        cum_sum += dd
        if cum_sum < 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
# => array([1.  , 1.75, 2.25, 2.5 , 2.5 , 2.25, 1.75, 1.  , 0.  ])

first example

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

# => array([0.  , 0.  , 0.  , 0.  , 0.  , 0.25, 0.75, 1.5 , 2.5 ])

enter image description here

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?

\$\endgroup\$
  • \$\begingroup\$ Can these graphs only go below zero at the beginning/end - like in the example? Can they go below zero half way through? \$\endgroup\$ – Peilonrayz Jun 14 at 15:01
  • \$\begingroup\$ I'm pretty sure you can iterate over a Numpy 1D array just as if it was a list. So converting it to a list seems unnecessary. \$\endgroup\$ – AlexV Jun 14 at 15:10
  • \$\begingroup\$ @Peilonrayz none of the graphs ever go below zero? Flows are always positive. The difference I calculate is explicitly made positive. \$\endgroup\$ – Seanny123 Jun 14 at 15:56
4
\$\begingroup\$

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:

>>> 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:

>>> 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:

>>> 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:

>>> 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:

>>> 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]
\$\endgroup\$
  • \$\begingroup\$ I'm a bit confused by your explanation. I understand subtract a value to get zero. However, hows does it follow that I need a running minimum? \$\endgroup\$ – Seanny123 Jun 14 at 15:54
  • 1
    \$\begingroup\$ Because you need to subtract that from ever value afterwards. Because if you subtract at any point in the cumsum you have to subtract every point after this. You also need to subtract different amounts for different positions, so the subtractions only get larger. \$\endgroup\$ – Peilonrayz Jun 14 at 15:58
2
\$\begingroup\$

There's two steps I think you could leave out:

  1. Casting the diff as a list, you can iterate over a numpy array just fine
  2. Checking if cum_sum is smaller than zero, you can check this using max()

    def pos_diff_cum_sum(flow_in: np.ndarray, flow_out: np.ndarray) -> np.ndarray:
    
        sums = []
        cum_sum = 0
        diff = flow_in - flow_out
    
        for dd in diff:
            cum_sum = max(cum_sum + dd, 0)
            sums.append(cum_sum)
    
        return np.array(sums)
    
\$\endgroup\$
  • \$\begingroup\$ Surprisingly, replacing the if with max makes the function slightly slower, according to my benchmarking. \$\endgroup\$ – Seanny123 Jun 14 at 15:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.