正确使用numpy.nditer？

Question

I'm trying to do an array operation with numpy.nditer , but don't get the expected result.

My code is

import numpy as np
arr1 = - np.random.random((2,2))
arr2 = np.random.random((2,2))
arr = np.zeros((2,2))
it = np.nditer([arr1, arr2, arr], [], [['readonly'], ['readonly'], ['writeonly']])
for a1, a2, a in it:
    a = a1 if -a1 < a2 else a2
print arr
print it.operands[2]

I'm getting all zero results in both arr and it.operands[2] , but I expected values from either arr1 or arr2 . What would be the correct way to assign values to arr in the iteration?

Answer 1

Doing a = in Python will simply rebind the local variable a ; it won't affect what a contains.

With nditer , the iteration variables a1 , a2 and a are actually 0-d arrays. Thus, to change a , use the (slightly odd) a[()] = syntax:

for a1, a2, a in it:
    a[()] = a1 if -a1 < a2 else a2

Note, though, that your whole code can be simplified greatly by using np.where :

import numpy as np
arr1 = - np.random.random((2,2))
arr2 = np.random.random((2,2))
arr = np.where(-arr1 < arr2, arr1, arr2)