Python函数在打开时从文件中读取可变长度的数据 block

Question

我有包含许多时间步长数据的数据文件，每个时间步长都格式化为如下 block :

TIMESTEP  PARTICLES
0.00500103 1262
ID  GROUP  VOLUME  MASS  PX  PY  PZ  VX  VY  VZ
651 0 5.23599e-07 0.000397935 -0.084626 -0.0347849 0.00188164 0 0 -1.04903
430 0 5.23599e-07 0.000397935 -0.0837742 -0.0442293 0.0121046 0 0 -1.04903
384 0 5.23599e-07 0.000397935 -0.0749234 -0.0395652 0.0143401 0 0 -1.04903
971 0 5.23599e-07 0.000397935 -0.0954931 -0.0159607 0.0100155 0 0 -1.04903
....

每个 block 包含 3 个标题行和多行与时间步相关的数据(第 2 行的 int)。与 block 关联的数据行数可以从 0 到 10 百万不等。每个 block 之间可能有一个空行，但有时会丢失。

我希望能够逐 block 读取文件，在读取 block 后处理数据 - 文件很大(通常超过 200GB)，并且一个时间步就可以轻松加载到内存中。

由于文件格式的原因，我认为编写一个读取 3 个标题行、读取实际数据然后返回一个漂亮的 numpy 数组进行数据处理的函数会非常容易。 我习惯了 MATLAB，您可以简单地按 block 读取，而不是在文件末尾。我不太确定如何使用 python 执行此操作。

我创建了以下函数来读取数据 block :

def readBlock(f):
    particleData = []
    Timestep = []
    numParticles = []
    linesProcessed = 0

    line = f.readline().strip()
    if line.startswith('TIMESTEP'): 

        timestepHeaders = line.strip()
        varData = f.readline().strip()
        headerStrings = f.readline().strip().split(' ')
        parts = varData.strip().split(' ')
        Timestep = float(parts[0])
        numParticles = int(parts[1])
        while linesProcessed < numParticles:
            particleData.append(tuple(f.readline().strip().split(' ')))
            linesProcessed += 1

        mydt = np.dtype([ ('ID',int), 
                     ('GROUP', int),
                     ('Vol', float),
                     ('Mass', float),
                     ('Px', float),
                     ('Py', float),
                     ('Pz', float),
                     ('Vx', float),
                     ('Vy', float),
                     ('Vz', float),
                     ] )

        particleData = np.array(particleData, dtype=mydt)

    return Timestep, numParticles, particleData

我尝试像这样运行该函数:

with open(fileOpenPath, 'r') as file:
    startWallTime = time.clock()

    Timestep, numParticles, particleData = readBlock(file)
    print(Timestep)

    ## Do processing stuff here 
    print("Timestep Processed")

    endWallTime = time.clock()

问题是这只从文件中读取第一个数据 block 并在那里停止 - 我不知道如何让它循环遍历文件直到它到达末尾并停止。

任何关于如何完成这项工作的建议都会很棒。我想我可以使用单行处理编写一种方法，并进行大量 if 检查来查看是否处于时间步的末尾，但简单的函数似乎更容易、更清晰。

Answer 1

您可以使用 numpy.genfromtxt 的 max_rows 参数:

with open("timesteps.dat", "rb") as f:
    while True:
        line = f.readline()
        if len(line) == 0:
            # End of file
            break
        # Skip blank lines
        while len(line.strip()) == 0:
            line = f.readline()
        line2_fields = f.readline().split()
        timestep = float(line2_fields[0])
        particles = int(line2_fields[1])
        data = np.genfromtxt(f, names=True, dtype=None, max_rows=particles)

        print("Timestep:", timestep)
        print("Particles:", particles)
        print("Data:")
        print(data)
        print()

这是一个示例文件:

TIMESTEP  PARTICLES
0.00500103    4
ID  GROUP  VOLUME  MASS  PX  PY  PZ  VX  VY  VZ
651 0 5.23599e-07 0.000397935 -0.084626 -0.0347849 0.00188164 0 0 -1.04903
430 0 5.23599e-07 0.000397935 -0.0837742 -0.0442293 0.0121046 0 0 -1.04903
384 0 5.23599e-07 0.000397935 -0.0749234 -0.0395652 0.0143401 0 0 -1.04903
971 0 5.23599e-07 0.000397935 -0.0954931 -0.0159607 0.0100155 0 0 -1.04903
TIMESTEP  PARTICLES
0.00500103    5
ID  GROUP  VOLUME  MASS  PX  PY  PZ  VX  VY  VZ
971 0 5.23599e-07 0.000397935 -0.0954931 -0.0159607 0.0100155 0 0 -1.04903
652 0 5.23599e-07 0.000397935 -0.084626 -0.0347849 0.00188164 0 0 -1.04903
431 0 5.23599e-07 0.000397935 -0.0837742 -0.0442293 0.0121046 0 0 -1.04903
385 0 5.23599e-07 0.000397935 -0.0749234 -0.0395652 0.0143401 0 0 -1.04903
972 0 5.23599e-07 0.000397935 -0.0954931 -0.0159607 0.0100155 0 0 -1.04903

TIMESTEP  PARTICLES
0.00500103    3
ID  GROUP  VOLUME  MASS  PX  PY  PZ  VX  VY  VZ
222 0 5.23599e-07 0.000397935 -0.0837742 -0.0442293 0.0121046 0 0 -1.04903
333 0 5.23599e-07 0.000397935 -0.0749234 -0.0395652 0.0143401 0 0 -1.04903
444 0 5.23599e-07 0.000397935 -0.0954931 -0.0159607 0.0100155 0 0 -1.04903

这是输出:

Timestep: 0.00500103
Particles: 4
Data:
[ (651, 0, 5.23599e-07, 0.000397935, -0.084626, -0.0347849, 0.00188164, 0, 0, -1.04903)
 (430, 0, 5.23599e-07, 0.000397935, -0.0837742, -0.0442293, 0.0121046, 0, 0, -1.04903)
 (384, 0, 5.23599e-07, 0.000397935, -0.0749234, -0.0395652, 0.0143401, 0, 0, -1.04903)
 (971, 0, 5.23599e-07, 0.000397935, -0.0954931, -0.0159607, 0.0100155, 0, 0, -1.04903)]

Timestep: 0.00500103
Particles: 5
Data:
[ (971, 0, 5.23599e-07, 0.000397935, -0.0954931, -0.0159607, 0.0100155, 0, 0, -1.04903)
 (652, 0, 5.23599e-07, 0.000397935, -0.084626, -0.0347849, 0.00188164, 0, 0, -1.04903)
 (431, 0, 5.23599e-07, 0.000397935, -0.0837742, -0.0442293, 0.0121046, 0, 0, -1.04903)
 (385, 0, 5.23599e-07, 0.000397935, -0.0749234, -0.0395652, 0.0143401, 0, 0, -1.04903)
 (972, 0, 5.23599e-07, 0.000397935, -0.0954931, -0.0159607, 0.0100155, 0, 0, -1.04903)]

Timestep: 0.00500103
Particles: 3
Data:
[ (222, 0, 5.23599e-07, 0.000397935, -0.0837742, -0.0442293, 0.0121046, 0, 0, -1.04903)
 (333, 0, 5.23599e-07, 0.000397935, -0.0749234, -0.0395652, 0.0143401, 0, 0, -1.04903)
 (444, 0, 5.23599e-07, 0.000397935, -0.0954931, -0.0159607, 0.0100155, 0, 0, -1.04903)]