CODEWITHSUNDEEP
pythonnumpycython
ibell
ibell

Reputation: 1080

How can I efficiently use numpy to carry out iterated summation

Definition of Problem

I can express the problem I am trying to solve like this:

Given A [Nx1] and n[Nx1] and x[Mx1], I want to carry out this operation

S = sum([A[i]*x**n[i] for i in range(len(n))])

with the use of numpy. I think I can do this using something like broadcasting in numpy, but I can't make sense of the numpy docs. Can someone help me figure out how to do this effciently in numpy?

I have a working cython solution for this problem below that is pretty fast and I wonder if I can do it more easily using numpy and avoid cython entirely.

Cython solution

Here is a working implementation for this problem using cython to demonstrate the problem:

cimport cython
import numpy as np

@cython.boundscheck(False)
cpdef  sum_function( double [:] A, double [:] x, double [:] n, double [:] out):
    cdef int i,j
    cdef int Nx = len(x)
    cdef int Nn = len(n)

    out[:] = 0

    for i in xrange(Nx):
        for j in range(Nn):
            out[i] += A[j]*x[i]**n[j]

Upvotes: 1

Views: 179

Answers (2)

Warren Weckesser
Warren Weckesser

Reputation: 114921

You can use np.sum(A * x.reshape(-1, 1)**n, axis=1):

In [40]: A
Out[40]: array([ 1.,  2., -1.,  3.])

In [41]: n
Out[41]: array([ 2.,  1.,  1.,  3.])

In [42]: x
Out[42]: array([  5.,  10.,   1.])

In [43]: S = sum([A[i]*x**n[i] for i in range(len(n))])

In [44]: S
Out[44]: array([  405.,  3110.,     5.])

In [45]: np.sum(A * x.reshape(-1, 1)**n, axis=1)
Out[45]: array([  405.,  3110.,     5.])

x.reshape(-1, 1) has shape (3,1), and n has shape (4,), so the (broadcast) result of x.reshape(-1, 1)**n has shape (3,4); column k contains x**n[k]. A has shape (4,), so A * x.reshape(-1, 1)**n has shape (3,4); element (i,j) of that product contains A[j]*x[i]**n[j]. The desired result is the sum of this array along axis=1.

Here's the cython version with the same data:

In [46]: out = np.zeros_like(x)

In [47]: sum_function(A, x, n, out)

In [48]: out
Out[48]: array([  405.,  3110.,     5.])

Upvotes: 3

Andrew Marshall
Andrew Marshall

Reputation: 1569

Broadcasting will only work if M=N, or if x is scalar. If M > N, just do

S = sum(a*x[:N]**n)

Similarly for M < N.

Upvotes: 0

Related Questions