How can we give a Numpy array custom slice attributes?

Viewed 50

I'm using Numpy arrays to represent different multi-dimensional physical quantities. For example, I might have a (6,) array (let's call it linear) where the first 3 elements represent 3D position of an object, and the last 3 elements represent 3D velocity of that object. I might also have defined a (7,) array (let's call it angular) where the first 4 elements represent a quaternion describing an objects rotation and the last 3 elements represent the angular velocity.

To access parts of these arrays I might do things like:

position = linear[:3]
angular_velocity = angular[-3:]

When it gets overwhelming to remember the slices (which it does when the arrays incorporate more information) I can do:

class Linear:
    p_slice = slice(0, 3)
    v_slice = slice(3, 6)

class Angular:
    q_slice = slice(0, 4)
    omega_slice = slice(4, 7)

position = linear[Linear.p_slice]
angular_velocity = angular[Angular.omega_slice]

This is a bit neater, but I still have to make sure I always know which type of variable I'm working with otherwise I might use the wrong type of slice like linear[Angular.q_slice] (whoops, I thought linear was an instance of my angular orientation/velocity arrays). So I'd rather make sure I'm disallowed from using the wrong type of slice with something like:

class Linear(np.ndarray):
    def p(self):
        return self[:3]

# make an instance of Linear

position = linear.p

I'm considering subclassing ndarray as I suggested above, but I'm wondering if there are better ways (and I'd like to avoid having to deal with the gotchas of subclassing ndarray). I definitely want to treat my objects as ndarrays though, so that means I don't want to wrap them in some other object type and have to access the ndarray as an instance property.

0 Answers
Related