NumPy Array Indexing

Question

Simple question here about indexing an array to get a subset of its values. Say I have a recarray which holds ages in one space, and corresponding values in another. I also have an array which is my desired subset of ages. Here is what I mean:

ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])

What I'm trying to do is something like this:

data.values[data.ages==desired_ages]

But, it's not working.

Answer 1

You want to create an subarray containing only the values whose indexes are in desired_ages.

Python doesn't have any syntax that directly corresponds to this, but list comprehensions can do a pretty nice job:

result = [value for index, value in enumerate(data.values) if index in desired_ages]

However, doing it this way results in Python scanning through desired_ages for each element in data.values, which is slow. If you could insert

desired_ages = set(desired_ages)

on the line before, this would improve performance. (You can determine if a value in is a set in constant time, regardless of the set's size.)

---

Complete Example

import numpy as np

ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])

result = [value for index, value in enumerate(data.values) if index in desired_ages]
print result

Output

[0.45852624094611272, 0.0099713014816563694, 0.26695859251958864, 0.10143425810157047, 0.93647796171383935]

Answer 2

I changed your example a little, shuffle the order of ages:

import numpy as np
np.random.seed(0)
ages = np.arange(3,103)
np.random.shuffle(ages)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([4, 16, 29, 80])

If all the elements of desired_ages are in data.ages, you can sort data by age field first, and then use searchsorted() to find all the index quickly:

data.sort(order="ages") # sort by ages
print data.values[np.searchsorted(data.ages, desired_ages)]

or you can use np.in1d the get a bool array and use it as index:

print data.values[np.in1d(data.ages, desired_ages)]