Sum elements in array with NumPy

Question

I am annotating a plot with matplotlib with this code

for position, force in np.nditer([positions, forces]):
    plt.annotate(
        "%.3g N" % force,
        xy=(position, 3),
        xycoords='data',
        xytext=(position, 2.5),
        textcoords='data',
        horizontalalignment='center',
        arrowprops=dict(arrowstyle="->")
    )

It works fine. But if I have elements at the same position, it will stack multiple arrows on each other, i.e. if I have positions = [1,1,4,4] and forces = [4,5,8,9], it will make two arrow at position 1 and two arrows at position 4, on top of each other. Instead, I want to sum the forces and only create one arrow at position 1 with force 4+5=9 and one arrow at position 4 with force 8+9=17.

How can I do this with Python and NumPy?

Edit

I guess it could be something like

import numpy as np

positions = np.array([1,1,4,4])
forces = np.array([4,5,8,9])

new_positions = np.unique(positions)
new_forces = np.zeros(new_positions.shape)

for position, force in np.nditer([positions, forces]):
    pass

I wonder if there's a better title for this question. To me "sum elements in array" just means np.sum(array) but that's clearly not what this is about. — Iguananaut
– Iguananaut, Commented May 6, 2015 at 18:07

Alan · Accepted Answer · 2015-05-07 01:54:43Z

3

I'm not sure numpy offers help. Here's a Python solution:

from collections import defaultdict
result = defaultdict(int)
for p,f in zip(positions,forces):
    result[p] += f

positions, forces = zip(*result.items())
print positions, forces

Edit: I'm not sure what "I have to do it with numpy" means, but

import numpy as np
positions = np.array([1,1,4,4])
forces = np.array([4,5,8,9])
up = np.unique(positions)
uf = np.fromiter((forces[positions == val].sum() for val in up), dtype=int)

print up, uf

edited May 7, 2015 at 1:54

answered May 6, 2015 at 18:02

Alan

9,65020 silver badges26 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Jamgreen Over a year ago

It's a good solution but I have to do it with numpy. Do you think your code can be written with numpy? I've tried to make some of the code in my edit

Collectives™ on Stack Overflow

Sum elements in array with NumPy

Edit

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

Edit

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related