Search code examples
pythonnumpypython-itertoolsgekko

numpy sum axis 1 in pure python


This may seem like a strange question, but how do you rewrite in pure python next line:

np.sum(three_dim_matrix, axis=1).cumsum(axis=1)

cumsum is supposed to be applied to a two-dimensional matrix, so the code for cumsum I could already find:

from itertools import accumulate
[list(accumulate(row)) for row in two_dim_matrix]

If you're really wondering why I don't use numpy, the problem is that optimizers for MINLP (such, GEKKO) don't support defining objective functions in numpy features


Example:

example = np.array([[[ 70,  110,  130],
                     [-50, -100, -200]],

                    [[300,  140,  120],
                     [300,  140,  120]],

                    [[ 400, 180, -240],
                     [1000, 320,  560]]])

first_step = np.sum(example, axis=1)
# [[  20   10  -70]
#  [ 600  280  240]
#  [1400  500  320]]

second_step = np.cumsum(first_step, axis=1)
# [[  20   30  -40]
#  [ 600  880 1120]
#  [1400 1900 2220]]

Solution

  • In the two steps your example shows, data being the input list:

    first_step = [list(map(sum, zip(*rows))) for rows in data]
    second_step = [list(accumulate(row)) for row in first_step]
    

    Or both steps combined (should be faster, as it doesn't build intermediate lists):

    both_steps = [list(accumulate(map(sum, zip(*rows)))) for rows in data]
    

    Try it online!