Как предотвратить добавление двух массивов путем трансляции в numpy?

У Numpy очень мощный механизм вещания. Он даже может добавлять массивы 1x2 и 2x1 без предупреждения. Мне не нравится такое поведение: с вероятностью 99%, такое добавление является следствием моей ошибки, и я хочу, чтобы исключение было выбрано. Вопрос: есть ли что-то вроде:

работает только тогда, когда A и B имеют точно ту же форму?

Ответы

Ответ 1

Вы можете определить подкласс ndarray, который проверяет форму результата после вычисления. Вычисление выполняется, и мы проверяем форму результата, если он не имеет ту же форму, что и операнд, возникает исключение:

import numpy as np

class NoBCArray(np.ndarray):

    def __new__(cls, input_array, info=None):
        obj = np.asarray(input_array).view(cls)
        return obj

    def __array_wrap__(self, out_arr, context=None):
        if self.shape != out_arr.shape:
            raise ValueError("Shape different")
        return np.ndarray.__array_wrap__(self, out_arr, context)

a = NoBCArray([[1, 2]])
b = NoBCArray([[3], [4]])

a + b # this will raise error

Если вы хотите проверить перед вычислением, вам потребуется обернуть __add__:

def check_shape(opfunc):
    def checkopfunc(self, arr):
        if self.shape != arr.shape:
            raise ValueError("Shape different before calculation")
        else:
            return opfunc(self, arr)
    return checkopfunc

class NoBCArray(np.ndarray):

    __add__ = check_shape(np.ndarray.__add__)