Skip to content

Instantly share code, notes, and snippets.

@theSage21
Last active Oct 12, 2018
Embed
What would you like to do?
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))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment