Search code examples
pythonarraysnumpymultidimensional-arraysum

How to stack summing vectors to numpy 3d array?


I have a 3d numpy array that looks like so:

>>> g
array([[[ 1.,  1.,  1.,  1.,  1.],
        [ 0.,  0.,  0.,  0.,  0.],
        [ 1.,  2.,  3.,  4.,  6.]],

       [[ 0.,  0.,  0.,  0.,  0.],
        [11., 22., 33., 44., 66.],
        [ 0.,  0.,  0.,  0.,  0.]]])
  1. I know I can calculate a sum along the first axis with gs = g.sum(axis=1) that will result in this array:
>>> gs
array([[ 2.,  3.,  4.,  5.,  7.],
       [11., 22., 33., 44., 66.]])

How do I stack this summing array to the original one as the forth vector in each of the two inside groups? The expected result would be:

>>> g
array([[[ 1.,  1.,  1.,  1.,  1.],
        [ 0.,  0.,  0.,  0.,  0.],
        [ 1.,  2.,  3.,  4.,  6.],
        [ 2.,  3.,  4.,  5.,  7.]],

       [[ 0.,  0.,  0.,  0.,  0.],
        [11., 22., 33., 44., 66.],
        [ 0.,  0.,  0.,  0.,  0.],
        [ 11., 22., 33., 44., 66.]]])
  1. And I have the same question about the summing array along the 0th dimension which is calculated with gss = g.sum(axis=0) and looks like so:
>>> gss
array([[ 1.,  1.,  1.,  1.,  1.],
       [11., 22., 33., 44., 66.],
       [ 1.,  2.,  3.,  4.,  6.]])

How do I stack it to the original array to get the result shown below?

>>> g
array([[[ 1.,  1.,  1.,  1.,  1.],
        [ 0.,  0.,  0.,  0.,  0.],
        [ 1.,  2.,  3.,  4.,  6.]],

       [[ 0.,  0.,  0.,  0.,  0.],
        [11., 22., 33., 44., 66.],
        [ 0.,  0.,  0.,  0.,  0.]],

       [[ 1.,  1.,  1.,  1.,  1.],
        [11., 22., 33., 44., 66.],
        [ 1.,  2.,  3.,  4.,  6.]]])

Solution

  • np.concatenate([g, g.sum(axis=1, keepdims=True)], axis=1)
    

    Equivalent for axis=0