Search code examples
pythonnumpysum

Python: general sum over numpy rows


I want to sum all the lines of one matrix hence, if I have a n x 2 matrix, the result should be a 1 x 2 vector with all rows summed. I can do something like that with np.sum( arg, axis=1 ) but I get an error if I supply a vector as argument. Is there any more general sum function which doesn't throw an error when a vector is supplied? Note: This was never a problem in MATLAB.

Background: I wrote a function which calculates some stuff and sums over all rows of the matrix. Depending on the number of inputs, the matrix has a different number of rows and the number of rows is >= 1


Solution

  • According to numpy.sum documentation, you cannot specify axis=1 for vectors as you would get a numpy AxisError saying axis 1 is out of bounds for array of dimension 1.

    A possible workaround could be, for example, writing a dedicated function that checks the size before performing the sum. Please find below a possible implementation:

    import numpy as np
    
    M = np.array([[1, 4],
                 [2, 3]])
    
    v = np.array([1, 4])
    
    def sum_over_columns(input_arr):
        if len(input_arr.shape) > 1:
            return input_arr.sum(axis=1)
        return input_arr.sum()
    
    print(sum_over_columns(M))
    print(sum_over_columns(v))
    

    In a more pythonic way (not necessarily more readable):

    def oneliner_sum(input_arr):
        return input_arr.sum(axis=(1 if len(input_arr.shape) > 1 else None))