CODEWITHSUNDEEP
pythonarraysnumpy
Jacob Garber
Jacob Garber

Reputation: 364

Vectorize Loop Using Numpy

I have two 1D Numpy arrays index_data and charge, which are the same length and contain ints and floats respectively. I am trying to make a total_charge array as follows:

total_charge = []
for i in range(len(index_data)):
    if index_data[i] == 0:
        total_charge.append(charge[i])
    else:
        total_charge[-1] += charge[i]
total_charge = np.array(total_charge)

How would I vectorize this? Help me Numpy wizards, you're my only hope.

Upvotes: 1

Views: 109

Answers (1)

Warren Weckesser
Warren Weckesser

Reputation: 114976

A concrete example of index_data and charge would help clarify the expected input. From reading what you have, however, I think using np.add.reduceat can be used.

Apparently each new group of charges is indicated by a 0 in index_data; otherwise, the values in index_data are not used. total_charge is the sum of the charges in each group.

Here's an example of how reduceat could be used.

First, some demo data:

In [105]: index_data
Out[105]: array([0, 1, 1, 0, 1, 1, 1, 0, 1])

In [106]: charge
Out[106]: array([ 1.5,  2. ,  3. ,  2.5,  1.5,  1. ,  1. ,  2. ,  1. ])

zerolocs gives the indices where 0 occurs in index_data:

In [107]: zerolocs = where(index_data==0)[0]

In [108]: zerolocs
Out[108]: array([0, 3, 7])

Use np.add.reduceat to sum the charge groups.

In [109]: total_charge = np.add.reduceat(charge, zerolocs)

In [110]: total_charge
Out[110]: array([ 6.5,  6. ,  3. ])

Upvotes: 4

Related Questions