Dask read_csv: 跳过周期性出现的行

Dask read_csv: skip periodically ocurring lines

我想使用 Dask 在多个时间步读取一个大的原子坐标文件。格式称为 XYZ 文件,如下所示:

3 timestep 1 C 9.5464696279 5.2523477968 4.4976072664 C 10.6455075132 6.0351186102 4.0196547961 C 10.2970471574 7.3880736108 3.6390228968 3 timestep 2 C 9.5464696279 5.2523477968 4.4976072664 C 10.6455075132 6.0351186102 4.0196547961 C 10.2970471574 7.3880736108 3.6390228968

第一行是原子序号,第二行只是注释。 之后,列出原子及其名称和位置。 列出所有原子后,对下一个时间步重复相同的操作。

我现在想通过dask.dataframe.read_csv加载这样的轨迹。 但是,我不知道如何跳过包含原子序号和注释的周期性出现的行。这真的可能吗?

编辑:

可以通过以下方式将此格式读入 Pandas 数据帧:

atom_nr = 3

def skip(line_nr):
    return line_nr % (atom_nr + 2) < 2


pd.read_csv(xyz_filename, skiprows=skip, delim_whitespace=True,
            header=None)

但看起来 Dask 数据框不支持将函数传递给 skiprows。

编辑 2: MRocklin 的回答有效!为了完整起见,我写下了我使用的完整代码。

from io import BytesIO

import pandas as pd
import dask.bytes
import dask.dataframe
import dask.delayed


atom_nr = ...
filename = ...


def skip(line_nr):
    return line_nr % (atom_nr + 2) < 2


def pandaread(data_in_bytes):
    pseudo_file = BytesIO(data_in_bytes[0])
    return pd.read_csv(pseudo_file, skiprows=skip, delim_whitespace=True,
                       header=None)

bts = dask.bytes.read_bytes(filename, delimiter=f"{atom_nr}\ntimestep".encode())
dfs = dask.delayed(pandaread)(bts)

sol = dask.dataframe.from_delayed(dfs)
sol.compute()

唯一剩下的问题是:如何告诉 dask 只计算前 n 帧?目前看来完整的轨迹已经读完了。

简答

不,pandas.read_csv 和 dask.dataframe.read_csv 都没有提供这种功能(据我所知)

长答案

如果您可以编写代码将其中一些数据转换为 pandas 数据帧,那么您可以使用

一般来说,这可能类似于以下内容:

values = read_bytes('filenames.*.txt', delimiter='...', blocksize=2**27)
dfs = [dask.delayed(load_pandas_from_bytes)(v) for v in values]
df = dd.from_delayed(dfs)

每个 dfs 对应于大约 blocksize 字节的数据(然后直到下一个分隔符)。您可以控制您希望分区使用此块大小的精细程度。如果你愿意,你也可以 select 只有其中的几个 dfs 对象来获得更小部分的数据

dfs = dfs[:5]  # only the first five blocks of `blocksize` data