CODEWITHSUNDEEP
pythonnumpylist-comprehension
N. Dryas
N. Dryas

Reputation: 23

List comprehension with nested for loop, conditional, and accumulator

I am trying to convert this piece of code into a list comprehension:

a = np.random.rand(10) #input vector
n = len(a) # element count of input vector
b = np.random.rand(3) #coefficient vector
nb = len(b) #element count of coefficients
d = nb #decimation factor (could be any integer < len(a))
 
c = []
for i in range(0, n, d):
    psum = 0
    for j in range(nb):
        if i + j < n:
            psum += a[i + j]*b[j]
    c.append(psum)

I've tried following suggestions from:

For example:

from itertools import accumulate
c = [accumulate([a[i + j] * b[j] for j in range(nb) if i + j < n] ) for i in range(0, n, d)]

Later, when trying to get values from c (e.g. c[:index]):

TypeError: 'NoneType' object is not subscriptable

Or:

from functools import partial
def get_val(a, b, i, j, n):
    if i + j < n:
        return(a[i + j] * b[j])
    else:
        return(0)
c = [
         list(map(partial(get_val, i=i, j=j, n=n), a, b)) 
             for i in range(0, n, d) 
             for j in range(nb)
    ]

in get_val, return(a[i + j] * b[j])

IndexError: invalid index to scalar variable.

Or:

psum_pieces = [[a[i + j] * b[j] if i + j < n else 0 for j in range(nb)] for i in range(0, n, d)]
c = [sum(psum) for psum in psum_pieces]

As well as many other iterations of these approaches. Any guidance would be much appreciated.

Upvotes: 2

Views: 388

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114330

You really don't need to be using a list comprehension here. With numpy, you can create a fast pipelined solution that does not run any loops directly in the interpreter.

First convert a into a 2D array shaped (n // d, nb). The missing elements (i.e., where i + j >= n in the loop) can be zero since that will make the corresponding increment to psum zero:

# pre-compute i+j as a 2D array
indices = np.arange(nb) + np.arange(0, n, d)[:, None]
# we only want valid locations
mask = indices < n

t = np.zeros(indices.shape)
t[mask] = a[indices[mask]]

Now you can compute c directly as

(t * b).sum(axis=1)

I suspect that if you benchmark this solution against anything written in vanilla python not compiled with numba, it will be much faster.

Upvotes: 2

abc
abc

Reputation: 11929

If I've understood correctly what you want is something like

res = [sum(a[i+j]*b[j] for j in range(nb) if i+j < n) for i in range(0,n,d)]

For each i, this will add to the resulting list the sum of the products a[i+j]*b[j] for j that varies from 0 to nb-1 when i+j < n

Upvotes: 1

Related Questions