rnmhdn.py

import numpy as np
from tqdm import tqdm
from multiprocessing import Pool


file_name = "test.gro"  #Should be changed
frame_count = 14        #Should be changed
line_count = 4014	#Should be changed
Lz = 2.84               #Should be changed
Lx = 6.95
Ly = 6.632
dz = 0.02
zs = np.r_[0:Lz-dz:142j]
zs_len = len(zs)
all_mzs = np.zeros([frame_count, zs_len])
M = np.zeros(frame_count)
czs = np.zeros(zs_len)

    
def calc_m_M(frame, index):
    res = np.zeros(zs_len)
    for line in frame:
        res[int(line[3]/Lz*zs_len)] += line[0]
    res = np.add.accumulate(res)/(Lx*Ly*0.4238)*dz
    return res, index


frames = [(np.genfromtxt(file_name,
			 skip_header=2+line_count*frame_index+3*frame_index, max_rows = line_count,
			 converters={1:lambda s: (-2  if (str(s, "UTF-8").startswith("O")) else 1)},
			 usecols = np.r_[1,3:6]),
	   frame_index)
	   for frame_index in range(frame_count)]

with Pool() as pool:
    work = pool.imap_unordered(calc_m_M, frames)
    temp = Lx*Ly*dz
    for i, (out, frame_idx) in enumerate(tqdm(work, total=frame_count)):
        all_mzs[frame_idx, :] = out
        M[frame_idx] = temp*sum(out)
        

out = ''
for z in range(zs_len):
    mz_avg = (np.sum(all_mzs[:,z]))/(frame_count)
    M_avg = np.sum(M)/(frame_count)
    czs[z] = ((np.matmul(M,all_mzs)[z])/(frame_count)) - mz_avg * M_avg
    out += '{}\t{}\t{}\t{}\n'.format(zs[z], czs[z], mz_avg, M_avg)

C = Lx*Ly*np.sum(czs)*dz
with open("results.txt", 'w') as file:
    file.write("z\tc\t<mz>\t<M>\n{}\nC: {}\n".format(out,C))