Python科学计算(第2版)
上QQ阅读APP看书,第一时间看更新

2.2.4 广播

当使用ufunc函数对两个数组进行计算时,ufunc函数会对这两个数组的对应元素进行计算,因此它要求这两个数组的形状相同。如果形状不同,会进行如下广播(broadcasting)处理:

1)让所有输入数组都向其中维数最多的数组看齐,shape属性中不足的部分都通过在前面加1补齐。

2)输出数组的shape属性是输入数组的shape属性的各个轴上的最大值。

3)如果输入数组的某个轴的长度为1或与输出数组的对应轴的长度相同,这个数组能够用来计算,否则出错。

4)当输入数组的某个轴的长度为1时,沿着此轴运算时都用此轴上的第一组值。

上述4条规则理解起来可能比较费劲,下面让我们看一个实际的例子。

先创建一个二维数组a,其形状为(6,1):

    a = np.arange(0, 60, 10).reshape(-1, 1)
      a     a.shape
    ------  -------
    [[ 0],  (6, 1) 
     [10],         
     [20],         
     [30],         
     [40],         
     [50]]         

再创建一维数组b,其形状为(5,):

    b = np.arange(0, 5)
           b        b.shape
    ---------------  -------
    [0, 1, 2, 3, 4]  (5,)   

计算a与b的和,得到一个加法表,它相当于计算两个数组中所有元素对的和,得到一个形状为(6,5)的数组:

    c = a + b
              c             c.shape
    ----------------------  -------
    [[ 0,  1,  2,  3,  4],  (6, 5) 
     [10, 11, 12, 13, 14],         
     [20, 21, 22, 23, 24],         
     [30, 31, 32, 33, 34],         
     [40, 41, 42, 43, 44],         
     [50, 51, 52, 53, 54]]         

由于a和b的维数不同,根据规则1),需要让b的shape属性向a对齐,于是在b的shape属性前加1,补齐为(1,5)。相当于做了如下计算:

    b.shape = 1, 5
            b          b.shape
    -----------------  -------
    [[0, 1, 2, 3, 4]]  (1, 5) 

这样,加法运算的两个输入数组的shape属性分别为(6,1)和(1,5),根据规则2),输出数组的各个轴的长度为输入数组各个轴的长度的最大值,可知输出数组的shape属性为(6,5)。

由于b的第0轴的长度为1,而a的第0轴的长度为6,为了让它们在第0轴上能够相加,需要将b的第0轴的长度扩展为6,这相当于:

    b = b.repeat(6, axis=0)
            b         b.shape
    -----------------  -------
    [[0, 1, 2, 3, 4],  (6, 5) 
     [0, 1, 2, 3, 4],         
     [0, 1, 2, 3, 4],         
     [0, 1, 2, 3, 4],         
     [0, 1, 2, 3, 4],         
     [0, 1, 2, 3, 4]]         

这里的repeat()方法沿着axis参数指定的轴复制数组中各个元素的值。由于a的第1轴的长度为1,而b的第1轴的长度为5,为了让它们在第1轴上能够相加,需要将a的第1轴的长度扩展为5,这相当于:

    a = a.repeat(5, axis=1)
              a             a.shape
    ----------------------  -------
    [[ 0,  0,  0,  0,  0],  (6, 5) 
     [10, 10, 10, 10, 10],         
     [20, 20, 20, 20, 20],         
     [30, 30, 30, 30, 30],         
     [40, 40, 40, 40, 40],         
     [50, 50, 50, 50, 50]]

经过上述处理之后,a和b就可以按对应元素进行相加运算了。当然,在执行a + b运算时,NumPy内部并不会真正将长度为1的轴用repeat()进行扩展,这样太浪费内存空间了。由于这种广播计算很常用,因此NumPy提供了ogrid对象,用于创建广播运算用的数组。

    x, y = np.ogrid[:5, :5]
      x            y        
    -----  -----------------
    [[0],  [[0, 1, 2, 3, 4]]
     [1],                   
     [2],                   
     [3],                   
     [4]]                   

此外,NumPy还提供了mgrid对象,它的用法和ogrid对象类似,但是它所返回的是进行广播之后的数组:

    x, y = np.mgrid[:5, :5]
            x                  y        
    -----------------  -----------------
    [[0, 0, 0, 0, 0],  [[0, 1, 2, 3, 4],
     [1, 1, 1, 1, 1],   [0, 1, 2, 3, 4],
     [2, 2, 2, 2, 2],   [0, 1, 2, 3, 4],
     [3, 3, 3, 3, 3],   [0, 1, 2, 3, 4],
     [4, 4, 4, 4, 4]]   [0, 1, 2, 3, 4]]

ogrid是一个很有趣的对象,它像多维数组一样,用切片元组作为下标,返回的是一组可以用来广播计算的数组。其切片下标有两种形式:

●开始值:结束值:步长,和np.arange(开始值,结束值,步长)类似。

●开始值:结束值:长度j,当第三个参数为虚数时,它表示所返回的数组的长度,和np.linspace(开始值,,结束值,长度)类似。

    x, y = np.ogrid[:1:4j, :1:3j]
           x                  y          
    ---------------  --------------------
    [[ 0.        ],  [[ 0. ,  0.5,  1. ]]
     [ 0.33333333],                      
     [ 0.66666667],                      
     [ 1.        ]]                      

利用ogrid的返回值,我们很容易计算二元函数在等间距网格上的值。下面是绘制三维曲面(x,y)=xex2 -y2的程序:

    x, y = np.ogrid[-2:2:20j, -2:2:20j]
    z = x * np.exp( - x**2 - y**2)

图2-6为使用ogrid计算的三维曲面。

图2-6 使用ogrid计算二元函数的曲面

为了充分利用ufunc函数的广播功能,我们经常需要调整数组的形状,因此数组支持特殊的下标对象None,它表示在None对应的位置创建一个长度为1的新轴,例如对于一维数组a,a[None, :]和a.reshape(1, -1)等效,而a[:, None]和a.reshape(-1, 1)等效:

    a = np.arange(4)
      a[None, :]    a[:, None]
    --------------  ----------
    [[0, 1, 2, 3]][[0],     
    [1],     
    [2],     
    [3]]     

下面的例子利用None作为下标,实现广播运算:

    x = np.array([0, 1, 4, 10])    
    y = np.array([2, 3, 8])   
    x[None, :] + y[:, None]
    array([[ 2,  3,  6, 12],
           [ 3,  4,  7, 13],
           [ 8,  9, 12, 18]])

还可以使用ix_()将两个一维数组转换成可广播的二维数组:

    gy, gx = np.ix_(y, x)
            gx            gy        gx + gy      
    ------------------  -----  ------------------
    [[ 0,  1,  4, 10]]  [[2],  [[ 2,  3,  6, 12],
    [3],  [ 3,  4,  7, 13],
    [8]] [ 8,  9, 12, 18]]

在上面的例子中,通过ix_()将数组x和y转换成能进行广播运算的二维数组。注意数组y对应广播运算结果中的第0轴,而数组x与第1轴对应。ix_()的参数可以是N个一维数组,它将这些数组转换成N维空间中可广播的N维数组。