How to prevent adding two arrays by broadcasting in numpy?

前端 未结 1 1870
北海茫月
北海茫月 2021-01-04 04:10

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 addi

相关标签:
1条回答
  • 2021-01-04 04:41

    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__)
    
    0 讨论(0)
提交回复
热议问题