美文网首页Python与大气科学
Python气象数据处理与绘图(17):加快循环的运算速度

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

作者: 摸鱼咯 | 来源:发表于2020-03-23 20:28 被阅读0次

    最近比较忙好几天没有更新,今天就更新一个简单且实用的数据处理方法吧。
    在使用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

    相关文章

      网友评论

        本文标题:Python气象数据处理与绘图(17):加快循环的运算速度

        本文链接:https://www.haomeiwen.com/subject/trkgyhtx.html