How to overide NumPy ndarray attributes / properties?

Question

I have a class that has real and imag attributes,

class A():
    def __init__(self, real, imag):
        self.real = real
        self.imag = imag
        
    def __repr__(self):
        return repr((self.real, self.imag))
        
    def conjugate(self):
        return A(self.real, -self.imag)

It behaves like a complex number (but in my case, I want to realize a complex associative algebra element (e.g. a quantum operator), which is more general than a complex number, so I can not use the build-in complex type.)

A(1,2).real  # returns 1
A(1,2).conjugate() # returns (1,-2)

However, when I put it in a NumPy array, the conjugate function still works as expected, but not the real attribute.

import numpy

numpy.array([A(1,2)]).real
# returns array([(1, 2)], dtype=object)
# expect: array(1, dtype=object)

numpy.array([A(1,2)]).conjugate()
# returns array([(1, -2)], dtype=object)
# as expected

How to fix the class A to override NumPy's real and imag attributes when the class object is in the numpy.ndarray?

Answer 1

Your class, with repr tweaked:

In [181]: class A:
     ...:     def __init__(self, real, imag):
     ...:         self.real = real
     ...:         self.imag = imag
     ...: 
     ...:     def __repr__(self):
     ...:         return f"A: {repr((self.real, self.imag))}"
     ...: 
     ...:     def conjugate(self):
     ...:         return A(self.real, -self.imag)

make an instance:

In [192]: x = A(1, 2)
In [193]: x
Out[193]: A: (1, 2)
In [194]: x.real
Out[194]: 1
In [195]: x.conjugate()
Out[195]: A: (1, -2)

Put it in an array along with a couple of other values. The result is object dtype:

In [196]: arr = np.array([x, 1, 1 + 4j])
In [197]: arr
Out[197]: array([A: (1, 2), 1, (1+4j)], dtype=object)

Elements of the array:

In [198]: arr[0]
Out[198]: A: (1, 2)

Array conjugate runs the conjugate method of each element:

In [199]: arr.conjugate()
Out[199]: array([A: (1, -2), 1, (1-4j)], dtype=object)

real is defined as a property, not a method. It treats your A just as it does the complex element:

In [200]: arr.real
Out[200]: array([A: (1, 2), 1, (1+4j)], dtype=object)
In [201]: arr.imag
Out[201]: array([0, 0, 0], dtype=object)

Elements themselves have this attribute:

In [202]: arr[2].imag
Out[202]: 4.0
In [203]: arr[0].imag
Out[203]: 2

With object dtype arrays, numpy tries to apply the "appropriate" method or operator to each element.

You didn't define add, so we get this error:

In [204]: arr + arr
Traceback (most recent call last):
  Input In [204] in <module>
    arr + arr
TypeError: unsupported operand type(s) for +: 'A' and 'A'

In [205]: arr[1:] + arr[1:]
Out[205]: array([2, (2+8j)], dtype=object)

In short, you cannot access element attributes as though they were array attributes. attributes are not the same as methods.

Answer 2

After reading the NumPy document on Subclassing ndarray, I figured out a possible solution, by defining a subclass of numpy.ndarray to host the class A objects. Then it is possible to override the definition of real and imag.

class As(numpy.ndarray):
    def __new__(cls, input_array):
        obj = numpy.asarray(input_array).view(cls)
        return obj
    
    @property
    def real(self):
        return numpy.frompyfunc(lambda x: x.real,1,1)(self)
    
    @property
    def imag(self):
        return numpy.frompyfunc(lambda x: x.imag,1,1)(self)

Now one can construct the object array and realize the expected behavior.

As([A(1,2)]).real
# As([1], dtype=object)

As([A(1,2)]).conjugate()
# As([(1, -2)], dtype=object)