python numpy mask 平均性能

Question

好的，经过一些搜索我似乎无法找到直接解决这个问题的 SO 问题。我研究了屏蔽数组，虽然它们看起来很酷，但我不确定它们是否是我需要的。

考虑 2 个 numpy 数组：

zone_data 是一个二维 numpy 数组，其中包含具有相同值的元素块。这是我的 'zones'.

value_data 是具有任意值的二维 numpy 数组（zone_data 的精确形状）。

我寻找一个与 zone_data/value_data 形状相同的 numpy 数组，用每个区域的平均值代替区域编号。

示例...采用 ascii 艺术形式。

zone_data（4 个不同的区域）：

1, 1, 2, 2
1, 1, 2, 2
3, 3, 4, 4
3, 4, 4, 4

value_data:

1, 2, 3, 6
3, 0, 2, 5
1, 1, 1, 0
2, 4, 2, 1

我的结果，称之为result_data:

1.5, 1.5, 4.0, 4.0
1.5, 1.5, 4.0, 4.0
2.0, 2.0, 1.0, 1.0
2.0, 2.0, 1.0, 1.0

这是我的代码。就给我一个完美的结果而言，它工作得很好。

result_data = np.zeros(zone_data.shape)
for i in np.unique(zone_data):
    result_data[zone_data == i] = np.mean(value_data[zone_data == i])

我的数组很大，我的代码片段需要几秒钟。我认为我有知识差距，没有找到任何有用的东西。循环方面需要委托给库或其他东西...aarg!

我寻求帮助以加快速度！ Python大神们，求大神赐教！

编辑 -- 添加基准脚本

import numpy as np
import time

zones = np.random.randint(1000, size=(2000,1000))
values = np.random.rand(2000,1000)

print 'start method 1:'
start_time = time.time()

result_data = np.zeros(zones.shape)
for i in np.unique(zones):
    result_data[zones == i] = np.mean(values[zones == i])

print 'done method 1 in %.2f seconds' % (time.time() - start_time)

print
print 'start method 2:'
start_time = time.time()

#your method here!

print 'done method 2 in %.2f seconds' % (time.time() - start_time)

我的输出：

start method 1:
done method 1 in 4.34 seconds

start method 2:
done method 2 in 0.00 seconds

Answer 1

我以为我在 scipy 的某个地方见过这个，但我找不到了。你看过了吗？

无论如何，您可以通过更改循环来获得第一个改进：

result_data = np.empty(zones.shape)  # minor speed gain
for label in np.unique(zones):
    mask = zones==label
    result_data[mask] = np.mean(values[mask])

这样您就不会不必要地进行两次布尔比较。这会稍微减少执行时间。

Answer 2

您可以使用 np.bincount:

count = np.bincount(zones.flat)
tot = np.bincount(zones.flat, weights=values.flat)
avg = tot/count
result_data2 = avg[zones]

这给了我

start method 1:
done method 1 in 3.13 seconds

start method 2:
done method 2 in 0.01 seconds
>>> 
>>> np.allclose(result_data, result_data2)
True

python numpy mask 平均性能

python numpy mask mean performance

python

arrays

performance

numpy

masking