Python气象数据处理与绘图(17)：加快循环的运算速度

最近比较忙好几天没有更新，今天就更新一个简单且实用的数据处理方法吧。
在使用python时大家会发现有时候程序运行的异常缓慢。数据读取环节是我们不能控制的，那么只有在数据处理环节，我们才能尽可能的简化它。让python计算加快的关键是向量化，就是对数据进行矩阵运算，而不是循环运算。
python语言本身就决定了它不能像C或者Fortran一样实行动态运算，而在气象科研中又常常遇到大量的循环和逻辑判断，这就让程序的运行异常缓慢，我提供两个解决办法，第一个就是前面提到的向量化，举一个简单的例子，这个例子很蠢，但是比较容易理解。

#向量运算：
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[3,2,1],[6,5,4],[9,8,7]])
c = a + b

#非向量运算：
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[3,2,1],[6,5,4],[9,8,7]])
c = np.zeros((3,3))
for i in range(3):
    for j in renge(3):
        c[i,j] = a[i,j] + b[i,j]

相信大家都不会用第二种去运算，但是我这里只是一个简单的例子，其实大家翻翻以前写过的循环，应该都会发现自己使用过非向量运算，最简单的比如说计算相关系数时：

from scipy.stats import pearsonr
r = np.zeros((a.shape[1],a.shape[2]))
p = np.zeros((a.shape[1],a.shape[2]))

for i in range(a.shape[1]):
    for j in range(a.shape[2]):
            r[i,j], p[i,j] = pearsonr(b,a[:,i,j])

其实这里远没有NCL的相关系数计算方便，但是我目前了解的python计算相关系数的函数均没有指定维度的语句，也就是说它们都是求两个序列的相关系数，当我们计算一个场与一个序列的相关系数时，只能通过循环实现。
那么，在这种情况下，就可以用到我要讲的第二种方法，使用Numba加速！
关于Numba是什么的问题，大家感兴趣可以百度深入了解，通俗地讲，就是一个可以提高python遍历运算性能的库。安装方法：

conda install numba

使用numba很简单，就是自己定义一个python函数，然后添加一个装饰器来提高性能。还是同样的计算相关系数的程序：

import numba as nb
#为相关系数计算函数块提供@jit装饰器
@nb.jit()
def r_caculate(b,a):
    r = np.zeros((a.shape[1],a.shape[2]))
    p = np.zeros((a.shape[1],a.shape[2]))
    for i in range(a.shape[1]):
        for j in range(a.shape[2]):
            r[i,j], p[i,j] = pearsonr(b, a[:,i,j])
    return r,p

f_hadsic =  xr.open_dataset('HadISST_ice.nc')
sic_had = np.array(f_hadsic['sic'])                     
r,p = r_caculate(pc, sic_had)

运算结果是4.39秒，而不适用numba加速时：

f_hadsic =  xr.open_dataset('HadISST_ice.nc')
sic_had = np.array(f_hadsic['sic'])                     
r,p = r_caculate(pc, sic_had)
r = np.zeros((sic_had.shape[1],sic_had.shape[2]))
p = np.zeros((sic_had.shape[1],sic_had.shape[2]))
for i in range(sic_had.shape[1]):
    for j in range(sic_had.shape[2]):
        r[i,j], p[i,j] = pearsonr(pc, sic_had[:,i,j])

计算花费了12.3秒。可见差距十分巨大。
当数据的循环和逻辑判断很大时，使用numba加速可以说会为我们节省大量的等待时间。
但是numba并不是万能的，据我目前使用的结果，它只支持numpy，scipy的一部分函数。具体还需要大家在使用过程中逐渐尝试。
再举一个例子(这个函数实现是在一个(时间，纬度，经度)数组中判断任意时刻任意一个点是否大于周围八个点的)：

@nb.jit()
def max_in_8surrounds(a):
    b = np.zeros((a.shape[0],a.shape[1],a.shape[2]))
    for t in range(a.shape[0]):
            for i in range(1,a.shape[1]-1):
                for j in range(1,a.shape[2]-1):
                    tmp = np.array([a[t,i-1,j-1],a[t,i+1,j+1],a[t,i-1,j+1],
                                   a[t,i+1,j],a[t,i-1,j],a[t,i,j+1],a[t,i,j-1],a[t,i+1,j-1]])
                    if(a[t,i,j]>np.max(tmp)):
                        b[t,i,j] = 1
    return b
point = max_in_8surrounds(a)

我计算的数组十分巨大，加入了修饰器后，比未使用numba加速可以说快了不止几十倍。其实这个程序还涉及了一个小技巧，就是我在判断一个点是否小于周围八个点时，我将周围八个点先同时放进了一个数组，然后用中心点去与数组的最大值去判断大小，这样远比进行八次逻辑判断的速度要快的多。

if ((a>a1) & (a>a2) & (a>a2) & (a>a2) & .................):

如果numba仍不能满足你的需要，那还是去学C或者Fortran吧QAQ