How to efficiently read an xyz-formatted text file in Julia?

Viewed 464

XYZ-formatted text files are quite common in chemistry. These can be very large plain text files which contain a sequence of the coordinates of many atoms. Here is an example of a two frames, but generically there might be, say, a million of them:

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

The goal is to read this into three different Arrays. It doesn't particulatly matter what types of arrays, but they should contain all of the headers, which consist of the number of atoms in the frame and the comment line. The next one should contain the atom labels for each frame, and finally all of the coordinates should be read into either an Array{Float64, 3} or Array{Arrray{Float64, 2}, 1. I prefer the latter, but if it's more efficient to use the former, then that is fine.

I have tried two three different approaches, and they all perform almost identically. One where I read the file line by line and parse each line into the appropriate arrays as I go, using push! to allocate new memory as I go. One where I read the entire file into a string and then parse this string, again, dynamically allocating the memory. Finally, I have tried to read the entire file into a string, then using the number of atoms in the first frame to pre-allocate all of the arrays with zeros or empty strings. Surprisingly to me, the pre-allocating approach is no faster, or maybe a bit slower, for a 650MB file containing 400,000 frames with 30 atoms each.

Here is my code for the version which first reads the file into a string and then parses this string.

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

And here is the timing of using this function on a 650MB file containing 400,000 frames each with 30 atoms.

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)

So, clearly, there is a lot of garbage collection going on which takes a whole bunch of time and is just really slow. The part that reads the file into a string seems to allocate twice as much memory as it should. I am unsure why that is.

I guess the next thing to try is to just read the file into a string and then parse the file in parallel chunks. I am hoping for some input that will at least help with the fact that my parsing seems to be resulting in lots of intermediate memory which has to get garbage collected, and I'm not actually sure how to fix this.

Any help in optimizing this function is greatly appreciated.

1 Answers

A significant source of allocations is I think this line (I have not checked it - it is just a guess):

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

you could rewrite it as:

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

and things should get faster (I have unrolled the loop as in this case it should be best I think, alternatively you could have used the loop, but the point is that parse.(Float64, coords[2:end]) allocates twice: 1) once with coords[2:end] which creates a copy and once with parse. which creates a temporary vector using broadcasting).

Also this line:

head = string(N)

can probably be removed as you do not use it later (but maybe compiler optimizes it out already).


Also as a general recommendation instead of:

for (i, line) in enumerate(file_contents)

I think it would be better to use while loop over i as a counter. The issue is that with your current loop you run isa(tryparse(Int, line), Int) for every line, and most of the time this is false while tryparse costs you (in other words - you process each line twice essentially while with while loop you would process each line only once).

Related