How to prevent adding two arrays by broadcasting in numpy?

Question

Numpy has a very powerful broadcasting mechanism. It can even add 1x2 and 2x1 array without any warning. I don't like such behaviour: with 99% percent probability such addition is a consequence of my error and I want an exception to be thrown. The question: is there something like:

numpy.safe_add(A,B)

that works only when A and B have exactly the same shape?

Answer 1

You can define a subclass of ndarray that checks the shape of the result after the calculation. The calculation is executed, and we check the shape of the result, if it's not the same shape as the operand, an exception is raised:

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

If you want to check before the calculation, you need wrap __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__)