CODEWITHSUNDEEP
pythonarraysnumpystructured-array
user2099143
user2099143

Reputation: 65

Working with NumPy structured arrays

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.

Upvotes: 0

Views: 145

Answers (1)

unutbu
unutbu

Reputation: 879541

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

Upvotes: 2

Related Questions