Numpy array of object attributes [duplicate]

Question

I have a multi-dimensional array of objects, something like:

 a = np.array([obj1,obj2,obj3])

The objects are instances of a class which has several attributes. Let's say one of them is heights and one of them is lengths. To get the corresponding multi-dimensional array of lengths and heights I do:

 lengths = np.array([obj1.length,obj2.length,obj3.length])

 heights = np.array([obj1.height,obj2.height,obj3.height])

This is starting to clutter up my code quite a lot. Is there a more efficient way of doing this? For instance, I had something like

 heights = a.height

in mind but obviously it doesn't work because a is an array of my objects and not my object. But is there something similar I can do that is efficient and pythonic? I tried something like

 for x in np.nditer(a,flags=['refs_ok']):
    print x.length

to see what would happen but it doesn't work because nditer returns a tuple somehow.

Any ideas?

Answer 1

You could vectorize the function:

>>> import numpy
>>> 
>>> class Obj(object):
...     def __init__(self, x, y):
...         self.x = x
...         self.y = y
... 
>>> arr = numpy.array([Obj(1, 2), Obj(3, 4), Obj(5, 6)])
>>> 
>>> vectorized_x = numpy.vectorize(lambda obj: obj.x)
>>> 
>>> vectorized_x(arr)
array([1, 3, 5])

Although I'm not sure if you should really store a NumPy array of Python objects in the first place. Vectorize is no more efficient than a Python loop. It would be more efficient to store an (n+1)-D array, as we could easily extract the contents simply by slicing which is a native operation, e.g.

>>> a = numpy.array([[(1, 2), (3, 4), (5, 6)], [(7, 8), (9, 10), (11, 12)], [(-13, -14), (-15, -16), (-17, -18)]])
>>> a[:,:,0]
array([[  1,   3,   5],
       [  7,   9,  11],
       [-13, -15, -17]])
>>> a[:,:,1]
array([[  2,   4,   6],
       [  8,  10,  12],
       [-14, -16, -18]])