python: Conditional plus operator for numpy array

Question

I want to apply the plus operator on two numpy arrays. However, there is a constriction. I want to sum up the element only if it is not zero in array one. I could do this in a loop, but this is very slow. Is there a numpy typical approach?

if a1[xyz] != 0:
    r[xyz] = a1[xyz] + a2[xyz]

Use masking : r = a1 + a2*(a1!=0), for zeros intialized r. — Divakar
– Divakar, Commented Dec 11, 2017 at 14:16
@Divakar OP's code seems to update an existing r not touching the values where a1==0 — Paul Panzer
– Paul Panzer, Commented Dec 11, 2017 at 14:19

Paul Panzer · Accepted Answer · 2017-12-11 14:17:39Z

3

Use a mask:

mask = a1 != 0
r[mask] = a1[mask] + a2[mask]

This assumes that r, a1, a2 have the same shapes.

answered Dec 11, 2017 at 14:17

Paul Panzer

53.3k3 gold badges59 silver badges103 bronze badges

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

2 Comments

dgrat Over a year ago

Thanks, weirdly I seems to break if I add more than one condition: mask = (img_array == 0 or img_array == 1) Do you know why?

Paul Panzer Over a year ago

or doesn't work element-wise. You must use |. Because of operator precedence you will then also have to put parentheses around the individual conditions.

B. M. · Accepted Answer · 2017-12-11 14:47:26Z

2

An other fast way : np.where(a1==0,r,a1+a2).

answered Dec 11, 2017 at 14:47

B. M.

18.7k2 gold badges40 silver badges56 bronze badges

Comments

hpaulj · Accepted Answer · 2017-12-11 18:39:30Z

As I noted in Optimize a function that acts on a numpy array with an if statement, many ufunc take a where parameter

In this case we can use np.add in such a way:

In [168]: r = np.arange(10)
In [169]: a1 = np.ones(10,int); a1[3:7]=0
In [170]: mask = a1.astype(bool)
In [171]: mask
Out[171]: array([ True,  True,  True, False, False, False, False,  True,  True,  True], dtype=bool)
In [172]: a2 = np.arange(10,20)

In [173]: a1+a2
Out[173]: array([11, 12, 13, 13, 14, 15, 16, 18, 19, 20])
In [174]: np.add(a1,a2)
Out[174]: array([11, 12, 13, 13, 14, 15, 16, 18, 19, 20])

In [175]: np.add(a1,a2, where=mask, out=r);
In [176]: r
Out[176]: array([11, 12, 13,  3,  4,  5,  6, 18, 19, 20])

Without the out, the where leaves 'random' values in the masked out elements.

In the previous post this where has about the same timing as the masked equivalent.

If you want to use a compound mask, try something like mask = (a1!=0) & (a2>15). The () are important.

Collectives™ on Stack Overflow

python: Conditional plus operator for numpy array

3 Answers 3

2 Comments

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

2 Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related