更高效的Python输入/输出

Question

我需要处理超过 1000 万个光谱数据集。数据结构如下:大约有 1000 个 .fits(.fits 是某种数据存储格式)文件，每个文件包含大约 600-1000 个光谱，其中每个光谱中有大约 4500 个元素(因此每个文件返回 ~1000 *4500 矩阵)。这意味着如果我要遍历 1000 万个条目，每个光谱将被重复读取大约 10 次(或者每个文件将被重复读取大约 10,000 次)。虽然相同的光谱被重复读取大约 10 次，但它不是重复的，因为每次我都提取相同光谱的不同部分。

我有一个目录文件，其中包含我需要的所有信息，例如坐标 x、y、半径 r、强度s 等。该目录还包含我要读取的目标文件(由 n1、n2 标识)和哪些光谱的信息在我将要使用的那个文件中(由 n3 标识)。

我现在的代码是:

import numpy as np
from itertools import izip
import fitsio

x = []
y = []
r = []
s = []
n1 = []
n2 = []
n3 = []
with open('spectra_ID.dat') as file_ID, open('catalog.txt') as file_c:
    for line1, line2 in izip(file_ID,file_c):
        parts1 = line1.split()
        parts2 = line2.split()
        n1.append(parts1[0])
        n2.append(parts1[1])
        n3.append(float(parts1[2]))
        x.append(float(parts2[0]))         
        y.append(float(parts2[1]))        
        r.append(float(parts2[2]))
        s.append(float(parts2[3]))  

def data_analysis(idx_start,idx_end):  #### loop over 10 million entries
    data_stru = np.zeros((idx_end-idx_start), dtype=[('spec','f4',(200)),('x','f8'),('y','f8'),('r','f8'),('s','f8')])

    for i in xrange(idx_start,idx_end)
        filename = "../../../data/" + str(n1[i]) + "/spPlate-" + str(n1[i]) + "-" + str(n2[i]) + ".fits"
        fits_spectra = fitsio.FITS(filename)
        fluxx = fits_spectra[0][n3[i]-1:n3[i],0:4000]  #### return a list of list
        flux = fluxx[0]
        hdu = fits_spectra[0].read_header()
        wave_start = hdu['CRVAL1']
        logwave = wave_start + 0.0001 * np.arange(4000)
        wavegrid = np.power(10,logwave)

    ##### After I read the flux and the wavegrid, then I can do my following analysis.

    ##### save data to data_stru

    ##### Reading is the most time-consuming part of this code, my later analysis is not time consuming.

问题是文件太大，没有足够的内存来一次加载它，而且我的目录没有结构化，所有将打开同一个文件的条目都分组在一起。我想知道是否有人可以提供一些想法将大循环分成两个循环:1)首先循环文件，这样我们就可以避免一次又一次地重复打开/读取文件，2)循环将要进入的条目使用相同的文件。

Answer 1

如果我正确理解您的代码，n1 和n2 确定打开哪个文件。那么，为什么不对它们进行 lexsort 呢？然后，您可以使用 itertools.groupby 对具有相同 n1、n2 的记录进行分组。这是缩小比例的概念证明:

import itertools

n1 = np.random.randint(0, 3, (10,))
n2 = np.random.randint(0, 3, (10,))
mockdata = np.arange(10)+100

s = np.lexsort((n2, n1))

for k, g in itertools.groupby(zip(s, n1[s], n2[s]), lambda x: x[1:]):
    # groupby groups the iterations i of its first argument
    # (zip(...) in this case) by the result of applying the
    # optional second argument (here lambda) to i.
    # Here we use the lambda expression to remove si from the
    # tuple (si, n1si, n2si) that zip produces because otherwise
    # equal (n1si, n2si) pairs would still be treated as different
    # because of the distinct si's. Hence no grouping would occur.
    # Putting si in there in the first place is necessary, so we
    # we can reference the other records of the corresponding row
    # in the inner loop.
    print(k)
    for si, n1s, ns2 in g:
        # si can be used to access the corresponding other records
        print (si, mockdata[si])

打印类似的东西:

(0, 1)
4 104
(0, 2)
0 100
2 102
6 106
(1, 0)
1 101
(2, 0)
8 108
9 109
(2, 1)
3 103
5 105
7 107

您可能希望在 lexsort 中包含 n3，但不包含分组，以便您可以按顺序处理文件的内容。