如何让你的矩阵运算速度提高4000+倍

在用Python进行矩阵运算（尤其是大型矩阵运算）的时候，最忌讳的是写循环，循环的执行效率极其的低，想要提高计算效率，有很多方法可以尝试，今天我们就来看一下如何在仅基于numpy的条件下，召唤一些技巧来加速矩阵的计算效率。

假如说有这样一道题：有一个中国区的海拔数据（DEM），是个二维矩阵，问：如何快速从中挑选出海拔高度大于等于4000米的点并将低于4000米的点赋值为0。

我们先来以正常循环的逻辑来解这道题，方法当然就是双层for循环，在每个点上判断值的大小是否大于等于4000，如果小于4000则将位置赋值为0，代码如下：

import copy
from cnmaps.sample import load_dem

lons, lats, dem = load_dem()
def myloop(dem):
    ndem = copy.deepcopy(dem)
    for i in range(dem.shape[0]):
        for j in range(dem.shape[1]):
            if ndem[i,j] < 4000:
                ndem[i,j] = 0
    return ndem

我们先来看看这个函数执行的耗时。

%timeit loop_dem = myloop(dem)1.25 s ± 43.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

本例中dem的维度是(225, 350)，for循环耗时1.25s。

下面我们来尝试一下用numpy的vectorize方法，将函数向量化。

vectorize函数向量化

vectorize是numpy的一个将函数向量化的方法，在官方文档中有专门的介绍。

定义一个向量化函数，该函数以嵌套的对象序列或 numpy 数组作为输入，并返回单个 numpy 数组或 numpy 数组的元组。向量化函数对输入数组的连续元组(如 python map 函数)计算 pyfunc，但它使用 numpy 的广播规则。

向量化输出的数据类型是通过使用输入的第一个元素调用该函数来确定的。这可以通过指定 otypes 参数来避免。

vectorize可以改造你的python函数，改造后的函数可以直接用于numpy向量运算风格的写法之中。

在官网的介绍中，还附加了这么一段描述：

提供向量化函数主要是为了方便，而不是为了性能，它执行的本质是一个for循环。

看到一句话，很多人就躺平了，觉得这玩意不会有性能上的提升，但 纸上得来终觉浅，绝知此事要躬行，实际上经过我的实验发现，使用vectorize向量化以后，相比于原生for循环在性能上是有非常显著提升的。还记得上面我们用原生for循环的成绩是1.25s吗？记住这个数字，下面看看vectorize能达到多少秒。

我们先来定义单次循环体的计算逻辑：

def myfilter(dem, threshold=4000):    
    if dem < threshold:        
        return 0    
    else:        
        return dem

然后把它向量化：

vfilter = np.vectorize(myfilter)

这个时候vfilter函数就是一个经过向量化改造的函数了，我们可以直接把numpy矩阵作为参数传进入进行矩阵运算： vector_dem = vfilter(dem)

我们来看看它的计算性能：

%timeit vector_dem = vfilter(dem)

结果是：

11.5 ms ± 206 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

11.5 ms vs 1250 ms (1.25s)，向量化以后的耗时仅为原生for循环的0.91%，速度提升了108倍！

这样看来，官网的那一番发言还是过于“谦虚”了。

另外细心的小伙伴可能发现了，np.vectorize本质上是一个闭包，所以其实你可以把它当成一个装饰器来使用，类似这样：

@np.vectorize
def myfilter(dem, threshold=4000):
    if dem < threshold:
        return 0
    else:
        return dem

然后调用的时候直接调用你定义的函数名即可：

vector_dem = myfilter(dem)

numpy的性能天花板

你以为np.vectorize就到头了吗？并不是，让我们来看看这个例子最快的实现——索引赋值：

ndem = copy.deepcopy(dem)%timeit ndem[ndem<4000] = 0

结果：

264 µs ± 8.23 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

WTF?!

264 us！又是一个量级上的提升！我们来把三次实验的单位统一一下：

• 原生for循环：1250000 us
• 向量化函数：11500 us
• 索引赋值：264 us

索引赋值的速度是向量化函数的43倍，是原生for循环的4734倍！

就我目前的经验来看，在不借助外力的情况下，召唤numpy性能天花板的方法应该是结合 花式索引 的各种骚操作。本质上矩阵运算的难点在于 逻辑分支，也就是在矩阵中实现类似于if-else的逻辑运算，只要你能在矩阵中实现了逻辑分支，任何分支内的运算步骤都可以使用矩阵运算轻易地实现。

这里所展示的只是一个最简单的例子，实际应用中，会有更复杂的场景，届时会非常考验开发者的思维水平和对numpy的熟练程度。

例如感兴趣的朋友可以细细品一下下面这段uv转风速风向的函数的实现，它可以直接传入矩阵形式的uv，快速计算出风速和风向，经过了长期的实战检验，可直接抄走使用：

def cal_wnswnd(u, v):
    """Calculate wind speed and wind direction by u, v components of wind

    Args:
        u (ndarray): U component of wind
        v (ndarray): V component of wind

    Returns:
        (ndarray, ndarray): Wind speed in m/s, wind direction in degree
    """
    ws = np.round(np.sqrt(u**2 + v**2), 1)

    wd = np.full_like(u, np.nan)  # wd = wind direction
    pvi = np.where(v>0)   # positive v index
    zvi = np.where(v==0)  # zero v index
    nvi = np.where(v<0)   # negative v index
    
    bzi = np.where((u==0)&(v==0))   # both uv zero index
    
    wd[pvi] = np.arctan(u[pvi]/v[pvi]) * 180/np.pi + 180
    wd[nvi] = np.arctan(u[nvi]/v[nvi]) * 180/np.pi
    wd[zvi] = np.where(u[zvi] > 0, 270, 90)
    wd[bzi] = np.full_like(wd[bzi], -9999)
    
    wd[wd<0] += 360
    wd = np.round(wd, 1)
    
    return ws, wd

前面说了这么多，最后我们来画张图，看一下实现的效果：

import matplotlib.pyplot as plt
plt.imshow(ndem[::-1])