使用NumPy结构化数组

Question

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.

Answer 1

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 ).