Избегайте numpy распределять операцию для перегруженного оператора

Question

По умолчанию numpy распределяет операции по массивам, если он не знает тип другого объекта. В большинстве случаев это хорошо работает. Например, следующее поведение ведет себя так, как ожидалось.

np.arange(5) + 5 # = [5, 6, 7, 8, 9] 

Я хотел бы определить класс, который переопределяет оператор сложения, как показано в приведенном ниже коде.

class Example: 
    def __init__(self, value): 
     self.value = value 

    def __add__(self, other): 
     return other + self.value 

    def __radd__(self, other): 
     return other + self.value 

Он хорошо работает для скалярных значений. Например,

np.arange(5) + Example(5) # = [5, 6, 7, 8, 9] 

Однако, это не совсем то, что я хочу для векторных значений. Например,

np.arange(5) + Example(np.arange(5)) 

дает выходной сигнал

array([array([0, 1, 2, 3, 4]), array([1, 2, 3, 4, 5]), 
    array([2, 3, 4, 5, 6]), array([3, 4, 5, 6, 7]), 
    array([4, 5, 6, 7, 8])], dtype=object) 

потому что __add__ оператор предыдущего Numpy массива имеет приоритет над __radd__ оператора, который я определил. Оператор Numbox __add__ вызывает __radd__ для каждого элемента массива numpy, дающего массив массивов. Как я могу избежать numpy, распределяющего операцию? Я хотел бы избежать подклассификации массивов numpy.

Answer 1

Для каждого np.ndarray и подклассов, которые не слишком хотят (например, в более ранних версиях Numpy np.ma.MaskedArray игнорировали его) вы можете определить __array_priority__, даже если вы не подкласс np.ndarray непосредственно.

Мысль об этом проста: подкласс с более высоким приоритетом определяет, какой оператор определяет математическую операцию, а не порядок операции.

рабочий пример с вами Example будет:

class Example: 

    # Define this priority 
    __array_priority__ = 2 

    def __init__(self, value): 
     self.value = value 

    def __add__(self, other): 
     return other + self.value 

    def __radd__(self, other): 
     return other + self.value 


import numpy as np 
np.arange(5) + Example(np.arange(5)) 
# returns array([0, 2, 4, 6, 8]) 

Так это работает, как хотелось. Но обратите внимание, что при использовании этого подхода есть некоторые тонкие проблемы:

Он не работает с MaskedArrays, потому что у них есть приоритет 15 (так что вам нужно будет изменить свой приоритет на 16+, чтобы он работал):

import numpy as np 
np.ma.array(np.arange(5)) + Example(np.arange(5)) 

# returns: 
masked_array(data = [array([0, 1, 2, 3, 4]) array([1, 2, 3, 4, 5]) array([2, 3, 4, 5, 6]) 
array([3, 4, 5, 6, 7]) array([4, 5, 6, 7, 8])], 
     mask = False, 
    fill_value = ?) 

и, например, он не работает с astropy.units.Quantity, потому что они определили их приоритет 10000:

import astropy.units as u 
(np.arange(5)*u.dimensionless_unscaled) + Example(np.arange(5)) 
#returns: 
<Quantity [array([ 0., 1., 2., 3., 4.]), 
      array([ 1., 2., 3., 4., 5.]), 
      array([ 2., 3., 4., 5., 6.]), 
      array([ 3., 4., 5., 6., 7.]), 
      array([ 4., 5., 6., 7., 8.])]> 

и это не работает с любым классом, который не использует numpy -машина.