Getting attributes from arrays of objects in NumPy

Question

Let's say I have an class called Star which has an attribute color. I can get color with star.color.

But what if I have a NumPy array of these Star objects. What is the preferred way of getting an array of the colors?

I can do it with

colors = np.array([s.color for s in stars])

But is this the best way to do it? Would be great if I could just do colors = star.color or colors = star->color etc like in some other languages. Is there an easy way of doing this in numpy?

Answer 1

The closest thing to what you want is to use a recarray instead of an ndarray of Python objects:

num_stars = 10
dtype = numpy.dtype([('x', float), ('y', float), ('colour', float)])
a = numpy.recarray(num_stars, dtype=dtype)
a.colour = numpy.arange(num_stars)
print a.colour

prints

[ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9.]

Using a NumPy array of Python objects usually is less efficient than using a plain list, while a recarray stores the data in a more efficient format.

Answer 2

You could use numpy.fromiter(s.color for s in stars) (note lack of square brackets). That will avoid creating the intermediate list, which I imagine you might care about if you are using numpy.

(Thanks to @SvenMarnach and @DSM for their corrections below).