Question

I’ve got a numpy array of custom objects. How can I get a new array containing the values of specific attributes of those objects?

Example:

import numpy as np

class Pos():
   def __init__(self, x, y):
      self.x = x
      self.y = y

arr = np.array( [ Pos(0,1), Pos(2,3), Pos(4,5) ] )

# Magic line
xy_arr = .... # arr[ [arr.x,arr.y] ]

print xy_arr
# array([[0,1],
         [2,3],
         [4,5]])

I should add that my motives for such an operation is to calculate the centre of mass of the objects in the array.

Answer 1

Usually, when I have multiple quantities that belong together and I want to benefit from numpys indexing power I use record arrays. Beware, if you do a lot of append/remove operations, numpy might be rather ineffective in terms of speed.

If I understood your comment correctly, this is an example where two values are selected by a third:

import numpy as np

# create a table for your data
dt = np.dtype([('A', np.double), ('x', np.double), ('y', np.double)])
table = np.array([(1,1,1), (2,2,2), (3,3,3)], dtype=dt)

# define a selection mask
selection = table['A'] > 1.5
columns = ['x', 'y']

print table[selection][columns]

A nice side effect is that saving this table using h5py is very simple and convenient as your data is already labeled.