numpy 在多维数组上Map函数

h7appiyu  于 2023-06-23  发布在  其他
关注(0)|答案(1)|浏览(213)

我有一个很大的维数数组MxNxK。我想通过M轴循环并在每个NxK数组上应用一个函数。我想避免一个for循环。
为了说明这一点,我构建了一个10 x5 x4的数组。

import numpy as np

data = []
for i in range(1,11):

    a =np.full(shape=(5,4),fill_value=i)
    data.append(a)

data = np.array(data)

现在,我遍历第一维,并在每个5x 4数组中应用一个简单的sum/mean。

means = []
sums = []

for i in data:
    means.append(i.sum())
    sums.append(i.mean())

实际上,我有一个5000 x200 x20数组,我想在每个200 x20数组上应用一个函数,但我想避免循环。“apply_沿着_axis”似乎没有解决这个问题。请帮帮忙

编辑回应@chrslg的回答,添加更多信息。在我的例子中,我想在MxNxK数组中的M个子数组中的每个子数组中应用prob_dist(pdf)和第N个百分位数。我注意到,使用numba teh第一次函数调用确实需要太长时间(7秒),但每次连续执行只需要0.77秒。而传统的M子阵循环需要3.9s

yhuiod9q

yhuiod9q1#

在这种情况下

# keeping the same mean/sum logic inversion as in your code
means = data.sum(axis=(1,2))
sums = data.means(axis=(1,2))

应该工作。
一般来说,有一些方法可以避免for循环对所有数据应用函数。例如apply along axis。或vectorize

def mymean(subarr):
    return subarr.mean()

mymeanvec = np.vectorize(mymean, signature='(m,n)->()')

mymeanvec(data)

是一种方法。
但这是不可取的。这些函数只是避免了for循环,而不是对python函数的M调用。这就是为什么numpy的文档明确指出这些函数不是为了性能。
真正矢量化的好方法是思考矢量化。大多数numpy操作可以直接在数组上工作,并自己执行for循环。
另一种方法是,当这不可能的时候,因为你正在做的事情太奇特了,没有for循环,使用numpy操作的组合来编写,那就是使用numba

from numba import jit

@jit(nopython=True)
def myNumbaMean(data):
    # It is always better to have the correctly shaped data ready, and avoid `append`. Especially with numpy (pure python append is efficient)
    means = np.empty((len(data),))
    sums = np.empty((len(data),))
    for i in range(len(data)):
        # still the same inversion
        means[i] = data[i].sum()
        sums[i] = data[i].mean()
    return means, sums

第一次调用需要一些时间,因为它会编译它。它和C代码一样快。
事实上,使用numba,如果你想做一些更奇特的事情,你甚至可以不依赖.mean().sum()。即使是这种看似幼稚的代码,也会相当高效

@jit(nopython=True)
def myNaiveMean(data):
    means = np.empty((len(data),))
    sums = np.zeros((len(data),))
    m,n,p=data.shape
    # this time without the inversion
    for i in range(m):
        for j in range(n):
            for k in range(p):
                sums[i] += data[i,j,k]
        means[i] = sums[i]/n/p
    return means,sums

计时

一些计时注意事项在数据形状(5000,200,20)上,计时如下所示:
| 方法|定时|
| - -----|- -----|
| python(yours)|133毫秒|
| 矢量化|119毫秒|
| 嫩巴|65毫秒|
| 直接numpy(我的第一个答案)|54毫秒|
| 嫩巴天真|35毫秒|
正如你所看到的,对于numba,这通常是相当令人困惑的,我们习惯于在python中认为最差的解决方案往往恰好是最好的解决方案。因为它只是在没有子调用和不必要的中间结果的情况下完成工作(就像numba版本中我们调用np.meannp.sum)。
而且,甚至令人讨厌的是,这个天真的版本甚至击败了纯粹的numpy版本(data.sum(axis=(1,2)))。话虽如此,我不建议使用numba当你有一个纯numpy版本。因为,的确,这一次,它战胜了纯粹的麻木;但这是因为纯numpy代码没有比我的代码更聪明(因为这里没有聪明的空间:它只是一个求和,你不能避免迭代所有元素并对它们求和。与您的代码相比,唯一使numpy更快的是它是编译的C代码与解释的python代码。但对numba来说,它是编译代码vs编译代码)。但是对于许多计算,numpy的作者已经实现了比我提出的第一个算法更有效的算法。因此,用numba重写所有内容不仅是浪费时间,而且往往导致函数甚至不如numpy快。
所以,一般来说,我总是试着“思考麻木”,就像我在第一个回答中所做的那样。然后,如果这是不可能的,写一个numba代码。然后,如果这是不可能的(例如,因为numba不可用),尝试“Map一个函数”,但要知道这只节省了for循环本身的时间(迭代),而不是for循环的内容,大部分时间,这是CPU被消耗的地方。
最后备注:在这种情况下,我的时间安排不是那么令人印象深刻。最多4倍增益。这似乎已经很多了。但这与我们在这里经常看到的抑制python的循环(其中×100或×1000增益因子非常常见)相比根本不算什么。这是因为你的幼稚代码并不那么慢。因为只有一个for循环,最外面的一个是用python做的。两个隐式内部循环位于summean代码中。总共有5000 + 5000×200 + 5000×200×20次迭代。5000×200 + 5000×200×20已经在numpy中矢量化了。所以我们在这里节省的只是5000次迭代和5000次调用。

相关问题