Python: general sum over numpy rows-CodePudding

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

CodePudding user response：

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

CodePudding user response：

You can do

np.sum(np.atleast_2d(x), axis=1)

This will first convert vectors to singleton-dimensional 2D matrices if necessary.