Add value to every "other" field ((i+j)%2==0) of numpy array

Question

I have an m-by-n numpy array, and I'd like to add 1.0 to all entries [i, j] when (i + j) % 2 == 0, i.e., "to every other square".

I could of course simply iterate over the fields

import numpy as np

a = np.random.rand(5, 4)

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        if (i + j) % 2 == 0:
            a[i, j] += 1.0

but needless to say this is really slow.

Any idea of how to improve on this?

Answer 1

You can easily do the operation in two steps, like

import numpy as np

a = np.zeros((5, 14))

# Even rows, odd columns
a[::2, 1::2] += 1
# Odd rows, even columns
a[1::2, ::2] += 1

print a

Answer 2

Here's one way using NumPy broadcasting -

a[(np.arange(a.shape[0])[:,None] + np.arange(a.shape[1]))%2==0] += 1

Explanation : We basically create two arrays that are equivalent of the i-th and j-th iterators. Let's call them I and J.

I = np.arange(a.shape[0])
J = np.arange(a.shape[1])

Now, to perform an operation between all possible i and j, we create extend I to 2D by pushing its elements into the first axis and thus creating singleton dimension along its second axis.

Figuratively, the effect of broadcasting could be put like this :

    I[:,None] : M , 1
            J : 1 , N
I[:,None] + J : M,  N

Thus, the final setting would be -

a[(I[:,None] + J)%2==0] += 1

To put it other way with the intention to avoid comparing with 0 and directly use mod-2 which would be essentially 0 or 1 -

a += (np.arange(a.shape[0])[:,None]-1 + np.arange(a.shape[1]))%2

---

One can also use np.ix_ to process odd and then even rows for setting, like so -

a[np.ix_(np.arange(0,a.shape[0],2),np.arange(0,a.shape[1],2))] += 1
a[np.ix_(np.arange(1,a.shape[0],2),np.arange(1,a.shape[1],2))] += 1

Answer 3

One can build a mask for "every other" element, and apply the addition on the mask.

# Create mask
m00 = np.zeros(a.shape[0], dtype=bool)
m00[0::2] = True
m01 = np.zeros(a.shape[1], dtype=bool)
m01[0::2] = True
m0 = np.logical_and.outer(m00, m01)

m10 = np.zeros(a.shape[0], dtype=bool)
m10[1::2] = True
m11 = np.zeros(a.shape[1], dtype=bool)
m11[1::2] = True
m1 = np.logical_and.outer(m10, m11)

m = np.logical_or(m0, m1)
a[m] += 1