如何用__array_ufunc__覆盖numpy ufunc

Question

I'm trying to implement numpy's ufunc to work with a class, using the __array_ufunc__ method introduced in numpy v1.13.

To simplify, here's what the class could look like :

class toto():
    def __init__(self, value, name):
        self.value = value
        self.name = name
    def __add__(self, other):
    """add values and concatenate names"""
        return toto(self.value + other.value, self.name + other.name)
    def __sub__(self, other):
    """sub values and concatenate names"""
        return toto(self.value - other.value, self.name + other.name)  

tata = toto(5, "first")  
titi = toto(1, "second")

Now if I try to apply np.add between these two, I get the expected result, as np.add relies on add . But if I call say np.exp, I get an error as expected :

>>> np.exp(tata)
AttributeError: 'toto' object has no attribute 'exp'

Now what I would like to do is to "override" all numpy ufuncs to work smoothly with this class without having to redefine every methods (exp(self), log(self), ...) in the class.

I was planning to use numpy ufunc's [__array_ufunc__] 1 to do this, but I don't really understand the doc as it doesn't provide a simple exemple of implementation.

If anyone has had any experience with this new functionnality that looks promising, could you provide a simple example ?

Answer 1

If I extend your class with a __array_ufunc__ method (and __repr__ ):

class toto():
    def __init__(self, value, name):
        self.value = value
        self.name = name
    def __add__(self, other):
        """add values and concatenate names"""
        return toto(self.value + other.value, self.name + other.name)
    def __sub__(self, other):
        """sub values and concatenate names"""
        return toto(self.value - other.value, self.name + other.name)

    def __repr__(self):
        return f"toto: {self.value}, {self.name}"
    def __array_ufunc__(self, *args, **kwargs):
        print(args)
        print(kwargs)

And try some ufunc calls:

In [458]: np.exp(tata)                                                          
(<ufunc 'exp'>, '__call__', toto: 5, first)
{}
In [459]: np.exp.reduce(tata)                                                   
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{}
In [460]: np.multiply.reduce(tata)                                              
(<ufunc 'multiply'>, 'reduce', toto: 5, first)
{}
In [461]: np.exp.reduce(tata,axes=(1,2))                                        
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{'axes': (1, 2)}
In [463]: np.exp.reduce(tata,axes=(1,2),out=np.arange(3))                       
(<ufunc 'exp'>, 'reduce', toto: 5, first)
{'axes': (1, 2), 'out': (array([0, 1, 2]),)}

That shows the information that your class receives. Evidently you can do what you want that. It can return NotImplemented . I suppose in your case it could apply the first argument to your self.value , or do some custom calculation.

For example if I add

      val = args[0].__call__(self.value) 
      return toto(val, self.name) 

I get:

In [468]: np.exp(tata)                                                          
(<ufunc 'exp'>, '__call__', toto: 5, first)
{}
Out[468]: toto: 148.4131591025766, first
In [469]: np.sin(tata)                                                          
(<ufunc 'sin'>, '__call__', toto: 5, first)
{}
Out[469]: toto: -0.9589242746631385, first

However if I put the object in an array, I still get the method error

In [492]: np.exp(np.array(tata))                                                
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-492-4dc37eb0ffe3> in <module>
----> 1 np.exp(np.array(tata))

AttributeError: 'toto' object has no attribute 'exp'

Apparently ufunc on an object dtype array iterates on the elements of the array, expecting to use a 'relevant' method. For np.add (+) it looks for the __add__ method. For np.exp it looks for an exp method. This __array_ufunc__ isn't called.

So it looks like it's intended more for a subclass of ndarray , or something equivalent. You, I think, are trying to implement a class that can work as elements of an object dtype array.