Чтение огромного файла xyz размером python (200M строк, ~ 15 ГБ) и выполнение над ним вычислений

Question

У меня есть выходной файл .xyz из программы молекулярного моделирования. Мне нужно рассчитать расстояние между двумя наборами атомов. Первая строка - это количество атомов (1046), следующую строку можно рассматривать как комментарий, который мне не нужен. Затем идут координаты атомов. Общий вид файла выглядит следующим образом.

    1046
 i =    13641, time =     5456.400, E =     -6478.1065220464
  O        -7.4658679231       -8.2711817669       -9.8539631371
  H        -7.6241360163       -9.2582538006       -9.9522290769
  H        -8.2358222851       -7.6941601822       -9.9653770757
  O        -4.9711266650       -4.7190696213      -15.2513827675
  H        -4.0601366272       -4.8452939622      -14.9451462873
  H        -5.3574156180       -5.6550789412      -15.1154558067
... ~ 1000 more lines
  O        -3.7163764338      -18.4917410571      -11.7137020838
  H        -3.3000068500      -18.5292231200      -12.6331415740
  H        -4.3493512803      -19.2443154891      -11.6751925772
    1046
 i =    13642, time =     5456.800, E =     -6478.1027892656
  O        -7.4669935102       -8.2716185134       -9.8549232159
  H        -7.6152044024       -9.2599276969       -9.9641510528
  H        -8.2364333010       -7.6943001983       -9.9565217204
  O        -4.9709831745       -4.7179801609      -15.2530422573
  H        -4.0670595153       -4.8459686871      -14.9472675802
  H        -5.3460565517       -5.6569802374      -15.1037050119
...

Индексы первого набора атомов, которые я хочу извлечь, имеют вид 0,3,6,9, ..., 360. Учитывая первые две строки заголовка в файле, я реализовал это как

w_index = list(range(2,362,3))

А другой набор состоит только из 8 атомов, которые я снова представил в виде списка.

c_index = [420,488,586,688,757,847,970,1031]

Я думаю, что нужно добавить соответствующие строки в отдельные списки, чтобы работать с ними с помощью функции dist_calculate_with_pbc_correction.

def open_file(filename):

    global step

    waters = []
    carbons = []
    step=0

    with open(filename, 'r') as infile:

        for index, line in enumerate(infile):
            items = line.split()

            if index % (natoms+2) in w_index:
                kind, x, y, z = items[0], float(items[1]), float(items[2]), float(items[3])
                waters.append([kind,x,y,z])

            if (index - 2)% (natoms+2) in c_index:
                kind, x, y, z = items[0], float(items[1]), float(items[2]), float(items[3])
                carbons.append([kind,x,y,z])

            if index > 0 and index % (natoms+2) == natoms:
                dist_calculate_with_pbc_correction(carbons,waters)
                carbons, waters = [], []
                step+=1

И это функция, которая выполняет все вычисления и записывает результаты в файл .

def dist_calculate_with_pbc_correction(c,w):

    write_file(output_fn,'%s %r \n' % ('#Step:',step))

    for i in range(len(c)):
        hydration = 0
        min_d = 100
        for j in range(len(w)):

            x_diff, y_diff, z_diff = c[i][1] - w[j][1], c[i][2] - w[j][2], c[i][3] - w[j][3]

            if x_diff > pbc_a/2:
                x_diff -= pbc_a
            elif x_diff < -pbc_a/2:
                x_diff += pbc_a

            if y_diff > pbc_b/2:
                y_diff -= pbc_b
            elif y_diff < -pbc_b/2:
                y_diff += pbc_b

            if z_diff > pbc_c/2:
                z_diff -= pbc_c
            elif z_diff < -pbc_c/2:
                z_diff += pbc_c

            dist = math.sqrt(x_diff**2 + y_diff**2 + z_diff**2)
            if dist < min_d:
                min_d, min_index = dist, w_index[j]-2
            if dist < r_cutoff :
                hydration +=1

        write_file(output_fn,'%d %s %d %d \n' % (c_index[i], round(min_d,3), min_index, hydration))


def write_file(filename,out):
    with open(filename,'a') as g:

        g.write(out)

Проблема в том, что этот код работает, но работает недостаточно быстро (читать очень медленно). Требуется около 27 минут для обработки go всего файла из ~ 200K шагов (~ 200M строк). Я говорю «очень медленно», потому что мой коллега придумал свою собственную версию на Фортране, которая делает то же самое, если не больше, работает в 10 раз быстрее. Ее код запускает весь файл менее чем за 3 минуты, хотя она пытается вычислить, чтобы определить наименьшее расстояние между всеми атомами, в отличие от того, как я диктую индексы атомов. Я знаю, что большую часть времени уходит на выбор атомов в функции open_file, чтобы добавить их в соответствующие списки, но я не знаю, как это улучшить. Любая помощь будет оценена по достоинству.

Вот образец файла углерода xyz, если вы хотите иметь под рукой образец файла.

import numpy as np
def sample_xyz(steps,natoms):

    with open('sample_xyz.xyz','w') as f:
        for step in range(steps):
            f.write(str(natoms)+'\n')
            f.write('i = {} \n'.format(step))

            for foo in range(natoms):
                line = 'C {} {} {}\n'.format(np.random.random()*10, np.random.random()*10, np.random.random()*10)
                f.write(line)