有没有办法增加不太慢的numpy数组？

Question

我有一个脚本，累积（计数）包含在两个文件中的字节。字节类似于C类unsigned char 0到255之间的整数值。

该累加器脚本的目标是计算这两个文件中字节的联合计数或频率。可能将其扩展到多个文件/维度。

这两个文件的大小相同，但它们非常大，约为6TB左右。

我正在使用numpy.uint64值，因为我使用Python的int类型出现溢出问题。

我有一个长度为255**2的1D累加器阵列，用于存储关节计数。

我计算从逐列到阵列偏移计算的偏移量，以便在正确的索引处增加联合频率。我以字节块（n_bytes）遍历这两个文件，将它们解压缩并增加频率计数器。

下面的代码的草图：

import numpy 
import ctypes 
import struct 

buckets_per_signal_type = 2**(ctypes.c_ubyte(1).value * 8) 
total_buckets = buckets_per_signal_type**2 
buckets = numpy.zeros((total_buckets,), dtype=numpy.uint64) 

# open file handles to two files (omitted for brevity...) 

# buffer size that is known ahead of time to be a divisible 
# unit of the original files 
# (for example, here, reading in 2.4e6 bytes per loop iteration) 
n_bytes = 2400000 

total_bytes = 0L 

# format used to unpack bytes 
struct_format = "=%dB" % (n_bytes) 

while True:  
    # read in n_bytes chunk from each file 
    first_file_bytes = first_file_handle.read(n_bytes) 
    second_file_bytes = second_file_handle.read(n_bytes) 

    # break if both file handles have nothing left to read 
    if len(first_file_bytes) == 0 and len(second_file_bytes) == 0: 
     break 

    # unpack actual bytes 
    first_bytes_unpacked = struct.unpack(struct_format, first_file_bytes) 
    second_bytes_unpacked = struct.unpack(struct_format, second_file_bytes) 

    for index in range(0, n_bytes): 
     first_byte = first_bytes_unpacked[index] 
     second_byte = second_bytes_unpacked[index] 
     offset = first_byte * buckets_per_signal_type + second_byte 
     buckets[offset] += 1 

    total_bytes += n_bytes 
    # repeat until both file handles are both EOF 

# print out joint frequency (omitted) 

与我曾经int的版本相比，这是令人难以置信的速度慢，一个数量级上的速度较慢。原来的作业在大约8个小时内完成（错误地，由于溢出），并且这个基于numpy的版本必须尽早退出，因为它似乎需要大约12-14天才能完成。

在这个基本任务中numpy的速度非常慢，或者我没有像Python那样用numpy做累加器。我怀疑后者，这就是我为什么要求帮助的原因。

我读了numpy.add.at，但我将添加到buckets数组中的解压字节数组没有偏移值，这些值自然地转换为buckets数组的“形状”。

是否有一种方法来存储和增加一个（长）整数数组，不溢出，哪个是合理的高性能？

我可以在C中重写这个，我猜，但我希望有一些numpy的东西，我忽略了这个问题很快就能解决。谢谢你的建议。

更新

我有旧版本numpy的和SciPy的那不支持numpy.add.at。所以这是另一个需要研究的问题。

我会尝试以下，看看如何继续下去：

first_byte_arr = np.array(first_bytes_unpacked)                     
second_byte_arr = np.array(second_bytes_unpacked)                     
offsets = first_byte_arr * buckets_per_signal_type + second_byte_arr                
np.add.at(buckets, offsets, 1L) 

希望它运行快一点！

更新II

使用np.add.at和np.array，这个工作将需要大约12天的完成。我现在要放弃numpy，并回到用C读取原始字节，其中运行时间更合理。谢谢你的建议！

Answer 1

不要试图按照所有的文件读取和struct代码，它看起来像你1添加到buckets阵列的各种插槽。这部分不应该花费几天时间。

但要想知道的dtype如何影响该步骤，我会测试将随机分类的索引加1。

In [57]: idx = np.random.randint(0,255**2,10000) 
In [58]: %%timeit buckets = np.zeros(255**2, dtype=np.int64) 
    ...: for i in idx: 
    ...: buckets[i] += 1 
    ...: 
9.38 ms ± 39.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
In [59]: %%timeit buckets = np.zeros(255**2, dtype=np.uint64) 
    ...: for i in idx: 
    ...: buckets[i] += 1 
    ...: 
71.7 ms ± 2.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each) 

uint64约慢了8倍。

如果没有重复，我们可以做buckets[idx] += 1。但允许重复，我们必须使用add.at：

In [60]: %%timeit buckets = np.zeros(255**2, dtype=np.int64) 
    ...: np.add.at(buckets, idx, 1) 
    ...: 
1.6 ms ± 348 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each) 
In [61]: %%timeit buckets = np.zeros(255**2, dtype=np.uint64) 
    ...: np.add.at(buckets, idx, 1) 
    ...: 
1.62 ms ± 15.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 

有趣的是D型细胞uint64不影响在这种情况下的时序。

你在评论中提到你尝试了一个列表累加器。我想这样：

In [62]: %%timeit buckets = [0]*(255**2) 
    ...: for i in idx: 
    ...: buckets[i] += 1 
    ...: 
3.59 ms ± 44.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 

这比数组的迭代版本更快。一般来说，数组上的迭代比列表上的迭代要慢。这是更快速的“全阵列”操作，例如add.at。

---

要验证add.at是迭代正确的替代品，比较

In [63]: buckets0 = np.zeros(255**2, dtype=np.int64) 
In [64]: for i in idx: buckets0[i] += 1 

In [66]: buckets01 = np.zeros(255**2, dtype=np.int64) 
In [67]: np.add.at(buckets01, idx, 1) 
In [68]: np.allclose(buckets0, buckets01) 
Out[68]: True 

In [69]: buckets02 = np.zeros(255**2, dtype=np.int64) 
In [70]: buckets02[idx] += 1 
In [71]: np.allclose(buckets0, buckets02) 
Out[71]: False 

In [75]: bucketslist = [0]*(255**2) 
In [76]: for i in idx: bucketslist[i] += 1 
In [77]: np.allclose(buckets0, bucketslist) 
Out[77]: True 

Answer 2

1. numpy都有自己的I文件fromfile/O方法，你可能会更好，如果使用你想要在numpy阵列中输出。 （见this question）

2. 可能更好地使用由numpy给出的array结构，使您buckets一个二维数组：

   buckets_per_signal_type = 2**(ctypes.c_ubyte(1).value * 8) 
   buckets = numpy.zeros((buckets_per_signal_type, buckets_per_signal_type), dtype=numpy.uint64) 
   

   ，然后只用np.add.at递增箱

   # define record_type to match your data 
   while True 
       data_1 = np.fromfile(first_file_handle, dtype=record_dtype, count=nbytes) 
       data_2 = np.fromfile(second_file_handle, dtype=record_dtype, count=nbytes) 
       s = np.minimum(data_1.size, data_2.size) 
       if s == 0: 
        break 
       np.add.at(buckets, [data_1[:s], data_2[:s]], 1)