并行向量化两个大型numpy数组的迭代

Question

我有两个很大的数组，类型是 numpy.core.memmap.memmap，分别叫 data 和 new_data，里面有超过700万个浮点数。

我需要在同一个循环中同时遍历这两个数组，目前我是这样做的：

for i in range(0,len(data)):
  if new_data[i] == 0: continue
  combo = ( data[i], new_data[i] )
  if not combo in new_values_map: new_values_map[combo] = available_values.pop()
  data[i] = new_values_map[combo]

不过这样做速度太慢了，所以我觉得应该使用numpy的向量化函数来提高效率。

我想知道是否可以用索引来进行向量化，这样向量化后的数组就可以将它的元素与另一个数组中对应的元素进行比较。

我考虑过把两个数组“打包”在一起，但我想这会造成不必要的开销。

有没有其他方法可以优化这个操作呢？

为了提供一些背景：我的目标是有效地合并这两个数组，使得两个数组中每一对对应值的独特组合在结果数组中用不同的值表示，除了在 new_data 数组中的零值会被忽略。这两个数组代表的是3D位图图像。

补充说明：available_values 是一组在 data 中尚未使用的值，并且在每次调用这个循环时都会保留。另一方面，new_values_map 在每次使用这个循环之前都会被重置为空字典。

补充说明2：数据数组只包含整数，也就是说：它最开始是初始化为零，然后在每次使用这个循环时，结合不同的 new_data，它会从 available_values 中填充更多的值，而 available_values 最初是一个整数范围。理论上，new_data 可以是任何东西。

Answer 1

Python提供了一些强大的工具来处理大量数据，比如生成器和迭代器。

简单来说，它们让你可以像使用普通列表一样访问数据，但不需要一次性把所有数据都加载到内存中，而是可以一块一块地取。

如果你需要同时访问两个大数组，你可以使用

for item_a, item_b in izip(data, new_data):
   #... do you stuff here

izip，它会创建一个迭代器，可以同时遍历你两个数组的元素，但它是根据需要一块一块地取，而不是一次性全部取出来。

Answer 2

看起来，把循环的前两行替换成下面的代码：

for i in numpy.where(new_data != 0)[0]:
  combo = ( data[i], new_data[i] )
  if not combo in new_values_map: new_values_map[combo] = available_values.pop()
  data[i] = new_values_map[combo]

就能达到想要的效果。

所以在这个循环中，大部分时间都是在遇到new_data中的零时跳过整个循环。我不太明白为什么这么多的空循环会消耗这么多资源，也许有一天我会明白的……

Answer 3

关于你问的向量化问题，答案可能是肯定的，不过你需要先弄清楚 available_values 里到底包含了什么，以及它是怎么用的，因为这正是向量化的关键。

你的解决方案可能看起来像这样……

indices = new_data != 0

data[indices] = available_values

在这种情况下，如果 available_values 可以看作是一组值，我们把第一组值分配给 data 中第一个 new_data 不为 0 的值，这样应该是可行的，前提是 available_values 是一个 numpy 数组。

假设 new_data 和 data 的值范围是 0 到 255，那么你可以构建一个 available_values 数组，里面包含 new_data 和 data 所有可能的唯一值对，像下面这样：

available_data = numpy.array(xrange(0, 255*255)).reshape((255, 255))
indices = new_data != 0
data[indices] = available_data[data[indices], new_data[indices]]

显然，available_data 可以是你想要的任何映射。上面的操作无论 available_data 里有什么，速度都应该很快（特别是如果你只构建一次 available_data 的话）。