ns_process_traj

#!/usr/bin/env python3

import argparse

p = argparse.ArgumentParser()
p.add_argument('-n','--no-headers', action='store_true', help="""process traj files with V and KE missing from headers""")
p.add_argument('-t','--trajectory', action='store_true', help="""produce combined trajectory""")
p.add_argument('-i','--interval', type=int, help="""interval to output trajectory""", default=1)
p.add_argument('-o','--output', help="""output file name""", default='stdout')
p.add_argument('-e','--energies_file', help="""energies file name for fixing restart-related inconsistencies""")
p.add_argument('traj_files', nargs='+', help="""trajectory files""")

args = p.parse_args()

import sys, re, subprocess
import numpy as np

if args.trajectory or args.no_headers:
   import ase.io

re_e =      re.compile('ns_energy=([^ ]*)')
re_volume = re.compile('volume=([^ ]*)')
re_iter = re.compile('iter=([^ ]*)')
re_KE =     re.compile('ns_KE=([^ ]*)')

def process_headers(file):
   proc = subprocess.Popen(['fgrep','Lattice',file],stdout=subprocess.PIPE)
   data = []
   for l in proc.stdout:

      l = str(l)
      r_e = re_e.search(l)
      if r_e is not None:
         r_e = float(r_e.group(1))

      r_volume = re_volume.search(l)
      if r_volume is not None:
         r_volume = float(r_volume.group(1))

      r_iter = re_iter.search(l)
      if r_iter is not None:
         r_iter = int(r_iter.group(1))

      r_KE = re_KE.search(l)
      if r_KE is not None:
         try:
            r_KE = float(r_KE.group(1))
         except:
            r_KE = 0.0

      data.append([r_iter, r_e, r_volume, r_KE])
   return data

def traj_iter(files, interval=1, energies_file=None):
    # print("traj_iter files interval Ts ", files, interval, Ts)
    n_files=len(files)
    n_done=0
    global_i = 0

    configs=[None]*n_files
    energies=np.zeros(n_files)
    atom_readers=[]
    # print("n_files ",n_files)
    # print("files ", files)
    for i in range(n_files):
        # print(i)
        # sys.stderr.write("starting file {}\n".format(files[i]))
        sys.stderr.write("%d" % (i % 10))
        sys.stderr.flush()
        ar = ase.io.iread(files[i])
        atom_readers.append(ar)
        configs[i] = next(ar)
        energies[i] = configs[i].info['ns_energy']
    sys.stderr.write("\n")
    sys.stderr.flush()

    if energies_file is not None:
        sys.stderr.write("starting energies file\n")
        with open(energies_file,"r") as f:
            top_line = f.readline()
            all_energies = np.loadtxt(f)

    sys.stderr.write("starting iterator\n")
    sys.stderr.flush()
    while True:
        # print("n_done ",n_done, "n_files", n_files)
        if n_done == n_files:
            # raise StopIteration()
            return

        # print("energies ", energies)
        highest_E_i = energies.argmax()
        energy_match=True
        if args.energies_file is not None:
            cur_iter = configs[highest_E_i].info['iter']
            E_ind_list = np.where(all_energies[:,0] == cur_iter)[0]
            ns_E = configs[highest_E_i].info['ns_energy']
            min_dE = np.min(np.abs(ns_E-all_energies[E_ind_list,1]))
            # print("check match",)
            # print("cur iter", cur_iter, "E_ind_list",E_ind_list)
            # print("ns_E",ns_E, "all_energies[E_ind_list]",all_energies[E_ind_list,1])
            # print("min_dE", min_dE, "ratio",min_dE/np.abs(ns_E))
            energy_match = (min_dE <= 1.0e-11 and np.abs(ns_E) < 1) or (min_dE/np.abs(ns_E) < 1.0e-11)
            # print("energy_match",energy_match)

        if energy_match:
            # print("highest_E_i ",highest_E_i)
            if global_i % interval == interval-1:
                if args.trajectory:
                    configs[highest_E_i].info['ns_global_i'] = global_i
                    yield configs[highest_E_i]
                else:
                    if configs[highest_E_i].has('masses') and configs[highest_E_i].has('momenta'):
                        KE = configs[highest_E_i].get_kinetic_energy()
                    else:
                        KE = None
                    x = [sum(configs[highest_E_i].numbers == Z) / len(configs[highest_E_i]) for Z in sorted(list(set(configs[highest_E_i].numbers)))]
                    yield [configs[highest_E_i].info['iter'], configs[highest_E_i].info['ns_energy'], configs[highest_E_i].get_volume(), KE] + [len(configs[highest_E_i])] + x

            global_i += 1

        try:
            configs[highest_E_i] = next(atom_readers[highest_E_i])
            energies[highest_E_i] = configs[highest_E_i].info['ns_energy']
            # print("incremented energies for stream ",highest_E_i,"to",energies[highest_E_i])
        except:
            # print("stream ",highest_E_i,"out of entries")
            n_done += 1
            configs[highest_E_i] = None
            energies[highest_E_i] = float('-inf')


if not args.trajectory:
   print("# i     H      V        KE")

if not args.trajectory and not args.no_headers:
    data=[]
    for file in args.traj_files:
        data.extend(process_headers(file))
    data.sort(key = lambda d: -d[1])
    for d in data[0::args.interval]:
        print(*d)
else:
   with open(args.output, 'w') as outfile:
       for d_i, d in enumerate(traj_iter(args.traj_files, args.interval, args.energies_file)):

          ibase = 100
          if d_i % (ibase*10) == ibase*10 - 1:
            sys.stderr.write(str((d_i // (ibase*10)) % 10)); sys.stderr.flush()
          elif d_i % ibase == ibase - 1:
            sys.stderr.write('.'); sys.stderr.flush()
          if d_i % (ibase*100) == ibase*100 - 1:
            sys.stderr.write('\n'); sys.stderr.flush()

          if args.trajectory:
             try:
                 ase.io.write(outfile, d, format=ase.io.formats.filetype(outfile, read=False))
             except:
                 ase.io.write(outfile, d, format='extxyz')
          else:
             print(*d)