如何在Julia中高效地读取xyz格式的文本文件



XYZ格式的文本文件在化学中很常见。这些可以是非常大的纯文本文件,其中包含许多原子的坐标序列。这里有一个两帧的例子,但一般来说可能有一百万个:

9
This is always a comment line
O -0.47895617997126 -2.66640293337835 -1.49681666046534
H  0.44977991701355 -2.60732102342602 -1.12187699450036
H -0.47021693081323 -3.17753067276769 -2.30844425585481
O  1.77621192889291 -0.57453626042269  2.55501382514084
H  1.89122710803544 -0.41488971167761  3.50755810540804
H 1.63216983845037 0.32650413720929 2.13442202830578
O  2.09888701915665 -2.00704728106228 -0.54677363666735
H  1.97657049717652 -1.12593867861129 -0.98496874845379
H  2.42388964325074 -1.74615001969454  0.33580032962549
[there might be new lines here but everything before the next frame is ignored]
9
Again, always a comment line
O -0.46951624407696 -2.67861471918262 -1.48616321880993
H  0.4349320578306  -2.59324562286125 -1.15158631652204
H -0.50124399561371 -3.21321032298656 -2.30894245367388
O  1.77946051665009 -0.57508897951565  2.54905851620706
H  1.84047894070088 -0.38790923183093  3.49683088026427
H 1.62214984757904 0.28130906816753 2.15144936348414
O  2.1037484479049  -2.00698121080077 -0.53624845433526
H  2.00667363126475 -1.10227312432539 -1.00287595161132
H  2.40066596189013 -1.75997079252843  0.34324093661981

目标是将其读取到三个不同的Arrays中。什么类型的数组并不重要,但它们应该包含所有的标头,这些标头由帧和注释行中的原子数组成。下一个应该包含每个帧的原子标签,最后所有坐标都应该读取到Array{Float64, 3}Array{Arrray{Float64, 2}, 1中。我更喜欢后者,但如果使用前者更有效,那也没关系。

我尝试了两三种不同的方法,它们的表现几乎完全相同。其中,我逐行读取文件,并在执行时将每一行解析为适当的数组,在执行时使用push!分配新内存。其中,我将整个文件读取为一个字符串,然后解析这个字符串,再次动态分配内存。最后,我尝试将整个文件读取为一个字符串,然后使用第一帧中的原子数来预分配所有带有零或空字符串的数组。令我惊讶的是,对于一个包含400000个帧(每个帧有30个原子(的650MB文件来说,预分配方法并不快,或者可能慢一点。

这是我的代码版本,它首先将文件读取为字符串,然后解析这个字符串。

function read_xyz(ifile::String)
"""
Reads in an xyz file of possibly multiple geometries, returning the header, atom labels, 
and coordinates as arrays of strings and Float64s for the coordinates.
"""
@time file_contents = readlines(ifile)
header = Array{String, 1}()
atom_labels = Array{Array{String, 1}, 1}()
geoms = Array{Array{Float64, 2}, 1}()
for (i, line) in enumerate(file_contents)
if isa(tryparse(Int, line), Int)
# allocate the geometry for this frame
N = parse(Int, line)
head = string(N)
labels = String[]
# store the header for this frame
head = string(line, file_contents[i+1])
i += 1
push!(header, head)
# loop through the geometry storing the vectors and atom labels as you go
geom = zeros((3, N))
for j = 1:N
coords = split(file_contents[i+1])
i += 1
push!(labels, coords[1])
geom[:,j] = parse.(Float64, coords[2:end])
end
push!(geoms, geom)
push!(atom_labels, labels)
end
end
return header, atom_labels, geoms
end

这是在一个650MB的文件上使用这个函数的时间,该文件包含400000个帧,每个帧有30个原子。

julia> @time header, labels, geoms = read_xyz("data/1body_forces.dat");
6.804485 seconds (25.22 M allocations: 1.226 GiB, 11.78% gc time)
33.436732 seconds (88.82 M allocations: 7.368 GiB, 23.17% gc time)
julia> @time header, labels, geoms = read_xyz("data/1body_forces.dat");
9.216689 seconds (25.22 M allocations: 1.226 GiB, 28.39% gc time)
36.765633 seconds (88.82 M allocations: 7.368 GiB, 29.41% gc time)

所以,很明显,有很多垃圾收集正在进行,这需要大量的时间,而且非常缓慢。将文件读取为字符串的部分似乎分配了两倍的内存。我不确定为什么会这样。

我想接下来要尝试的是将文件读取为字符串,然后将文件解析为并行块。我希望得到一些输入,这些输入至少有助于解决这样一个事实:我的解析似乎导致了大量的中间内存,这些内存必须被垃圾收集,而我实际上不知道如何解决这个问题。

非常感谢您对优化此功能的任何帮助。

分配的一个重要来源是我认为这一行(我没有检查它——这只是一个猜测(:

geom[:,j] = parse.(Float64, coords[2:end])

您可以将其重写为:

geom[1,j] = parse.(Float64, coords[2])
geom[2,j] = parse.(Float64, coords[3])
geom[3,j] = parse.(Float64, coords[4])

事情应该会变得更快(我已经展开了循环,因为在这种情况下,我认为这应该是最好的,或者你可以使用循环,但重点是parse.(Float64, coords[2:end])分配两次:1(一次是与创建副本的coords[2:end]分配,一次是使用广播创建临时向量的parse.分配(。

还有这一行:

head = string(N)

可能会被删除,因为以后不使用它(但编译器可能已经优化了它(。


也作为一般建议,而不是:

for (i, line) in enumerate(file_contents)

我认为最好使用while循环而不是i作为计数器。问题是,在当前循环中,您对每一行都运行isa(tryparse(Int, line), Int),大多数情况下这是错误的,而tryparse会让您付出代价(换句话说,您基本上处理每一行两次,而在while循环中,每一行只处理一次(。

最新更新