How to make 2D arrays with numpy while preserving inheritance

Question

I have a Polygon class that uses a 2d ndarray for storing it's points. Every point in the class is an instance of a Vector class, this Vector class inherits numpy's ndarray. Am currently using numpy.vstack to append points, but the problem is that my vector instances are converted to their parent class, numpy.ndarray.

import numpy as np


class Vector(np.ndarray):
    def __new__(cls, *args, **kwargs):
        arr = np.asarray([args[0], args[1], args[2]]).view(cls)
        arr = arr.astype('float64')
        return arr


class Polygon(object):
    def __init__(self):
        self.points = np.array([]).reshape(0, 3)

    def append_point(self, point):
        self.points = np.vstack([self.points, point])


poly = Polygon()
vec1 = Vector(1, 2, 3)
print(type(vec1))  # output: <class '__main__.Vector'>

poly.append_point(vec1)

print(poly.points)  # output: [[1. 2. 3.]]
print(type(poly.points[0]))  # output: <class 'numpy.ndarray'>

Is there a better approach for this?

Answer 1

I think you need to override the __array_function__ method if you want to preserve the subclass:

class Vector(np.ndarray):
    def __new__(cls, *args, **kwargs):
        arr = np.asarray([args[0], args[1], args[2]]).view(cls)
        arr = arr.astype('float64')
        return arr

    def __array_function__(self, func, types, args, kwargs):
         x,y = args[0]
         return func([x.view(np.ndarray), y.view(np.ndarray)], **kwargs).view(Vector)

This will give

poly = Polygon()
vec1 = Vector(1, 2, 3)    # Vector([1., 2., 3.])

poly.append_point(vec1)

poly.points    # Vector([[1., 2., 3.]])
poly.points[0]    # Vector([1., 2., 3.])

Any function you call using Vector will be intercepted by __array_function__ (except ufuncs). Although this works for your specific case you will need to modify it to make it more general for any kind of inputs.