I am working with a NumPy structured array with the following structure:
ar = np.array([(760., 0), (760.3, 0), (760.5, 0), (280.0, 1), (320.0, 1), (290.0, 1)], dtype=[('foo', 'f4'),('bar', 'i4')])
What is an efficient way of extracting the 'foo' fields for a specific value of 'bar'? For example, I would like to store all the 'foo' values for which 'bar' is 0 in an array:
fooAr = ar['foo'] if ar['bar'] is 0
The above does not work.
Use ar['foo'][ar['bar'] == 0]:
ar = np.array([(760., 0), (760.3, 0), (760.5, 0), (280.0, 1), (320.0, 1), (290.0, 1)], dtype=[('foo', 'f4'),('bar', 'i4')])
print(ar['bar'] == 0)
# array([ True, True, True, False, False, False], dtype=bool)
result = ar['foo'][ar['bar'] == 0]
print(result)
# array([ 760. , 760.29998779, 760.5 ], dtype=float32)
Note that since a boolean selection mask, ar['bar'] == 0, is used, result is a copy of parts of ar['foo'].
Thus, modifying result would not affect ar itself.
To modify ar assign to ar['foo'][mask] directly:
mask = (ar['bar'] == 0)
ar['foo'][mask] = 100
print(ar)
# array([(100.0, 0), (100.0, 0), (100.0, 0), (280.0, 1), (320.0, 1), (290.0, 1)],
# dtype=[('foo', '<f4'), ('bar', '<i4')])
Assignment to ar['foo'][mask] calls ar['foo'].__setitem__ which affects ar['foo'].
Since ar['foo'] is a view of ar, modifying ar['foo'] affects ar.
Note that the order of indexing matters here. If you tried applying the boolean mask
before selecting the 'foo' field, as in:
ar[mask]['foo'] = 99
Then this would not affect ar, since ar[mask] is a copy of ar.
Nothing done to the copy (ar[mask]) affects the original (ar).
