Getting Attributes from Arrays of Objects in NumPy

Question

Let's say I have a class called Star that has the color attribute. I can get the color with star.color .

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

I can do it with

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

But is this the best way to do this? 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 to do this in numpy?

Answer 1

The closest thing to what you want is to use recarray instead of ndarray 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 is usually less efficient than using a simple list , and recarray stores data in a more efficient format.

Answer 2

You can use numpy.fromiter(s.color for s in stars) (note the absence of square brackets). This will avoid creating an intermediate list, which, I think, you might need if you use numpy.

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