Newest version of lmp_ave_post as I use them on therm115

lammps/lmp_ave_post.py

@@ -9,6 +9,7 @@ and returns average and sem.
 import argparse
 import numpy as np
 import sys
+import os
 import logging
 import utilsscript
 
@@ -32,7 +33,7 @@ def get_fields(infile, is_non_standard):
     return fields
 
 
-def compute_mean(fields, infile):
+def compute_mean(fields, infile, stepmin, stepmax):
     '''
     Computes data ave and sem.
     '''
@@ -60,6 +61,11 @@ def compute_mean(fields, infile):
                 if not line:
                     break
                 step, nentries = int(line[0]), int(line[1])
+                keep = True
+                if stepmin and step < stepmin:
+                    keep = False
+                if stepmax and step > stepmax:
+                    break
                 if nrec % 1000 == 0:
                     logging.info(
                             "Reading record number {:<10d}, step {:<10.2f}."
@@ -72,10 +78,13 @@ def compute_mean(fields, infile):
                             )
                     sys.exit(1)
                 for i in range(nentries):
+                    if keep:
                         line = list(map(float, f.readline().split()))
                         for j, k in enumerate(fields):
                             ave[i, j] += line[j]
-                nrec += 1
+                    else:
+                        f.readline()
+                nrec += keep
             except EOFError:
                 break
 
@@ -92,6 +101,11 @@ def compute_mean(fields, infile):
             try:
                 line = line.split()
                 step, nentries = int(line[0]), int(line[1])
+                keep = True
+                if stepmin and step < stepmin:
+                    keep = False
+                if stepmax and step > stepmax:
+                    break
                 if nrec % 1000 == 0:
                     logging.info(
                             "Reading record number {:<10d}, step {:<10.2f}."
@@ -104,17 +118,19 @@ def compute_mean(fields, infile):
                             )
                     sys.exit(1)
                 for i in range(nentries):
+                    if keep:
                         line = list(map(float, f.readline().split()))
                         for j, k in enumerate(fields):
                             sem[i, j] += (line[j] - ave[i, j])**2
-                nrec += 1
+                    else:
+                        f.readline()
+                nrec += keep
                 line = f.readline()
                 if not line:
                     break
             except EOFError:
                 break
     sem = np.sqrt(sem/nrec)
-
     return ave, sem
 
 
@@ -131,8 +147,8 @@ def write_output(outfile, fields, ave, sem):
             ]
     logging.info("About to write output file.")
     with open(outfile, 'w') as f:
-        line = " ".join(('#', *["{:^17}"]*len(fields), '\n'))
-        f.write(line.format(*fields))
+        line = " ".join(('#', *["{:^17}"]*len(fields)))
+        f.write(''.join([line.format(*fields).rstrip(), '\n']))
         for i, (a, s) in enumerate(zip(ave, sem)):
             # I found no nice way of doing this...
             elems = []
@@ -141,9 +157,9 @@ def write_output(outfile, fields, ave, sem):
                     if fi == 'Chunk':
                         elems.append("{:^17d}".format(i+1))
                     else:
-                        elems.append("{:>8.3f} {:>8.3f}".format(a[j], 0.))
+                        elems.append("{:>12.6f} {:>12.6f}".format(a[j], 0.))
                 else:
-                    elems.append("{:>8.3f} {:>8.3f}".format(a[j], s[j]))
+                    elems.append("{:>12.6f} {:>12.6f}".format(a[j], s[j]))
             line = " ".join((*elems, '\n'))
             f.write(line)
     logging.info("DONE!")
@@ -177,6 +193,20 @@ def main():
         default="output.mean",
         help="Output file name.",
     )
+    parser.add_argument(
+        "--smin",
+        dest="stepmin",
+        default=0,
+        type=int,
+        help="Minimum step to consider [default %s]"
+        )
+    parser.add_argument(
+        "--smax",
+        dest="stepmax",
+        default=0,
+        type=int,
+        help="Maximum step to consider, 0 for all [default %s]"
+        )
     parser.add_argument(
         "-n",
         "--non-standard",
@@ -194,11 +224,17 @@ def main():
 
     infile = args.infile
     outfile = args.outfile
+    stepmin = args.stepmin
+    stepmax = args.stepmax
     is_non_standard = args.non_standard
 
+    if not os.path.isfile(infile):
+        print("Could not find file {}.".format(infile))
+        sys.exit()
+
     fields = get_fields(infile, is_non_standard)
 
-    ave, sem = compute_mean(fields, infile)
+    ave, sem = compute_mean(fields, infile, stepmin, stepmax)
 
     write_output(outfile, fields, ave, sem)
 

lammps/plot_forces.py

@@ -105,7 +105,7 @@ def main():
                 logging.info("Found {} token".format(tok))
                 infos = next(lines).split()
                 npoints = int(infos[1])
-                dummy = next(lines)
+                next(lines)
                 for i in range(npoints):
                     line = next(lines).split()
                     r.append(float(line[1]))
@@ -116,7 +116,7 @@ def main():
                 force = np.array(force)
                 toplot.append([r, ener, force, tok])
                 tok = []
-                dummy = next(lines)
+                next(lines)
             except StopIteration:
                 break
     for tp in toplot:
@@ -127,6 +127,8 @@ def main():
     for tp in toplot:
         axes[0].plot(tp[0], tp[1], label=tp[3], linewidth=3)
         axes[1].plot(tp[0], tp[2], label=tp[3], linewidth=3)
+        hmin, hmax = axes[1].get_xlim()
+        axes[1].hlines(0, hmin, hmax, color="black", linewidth=3, linestyles="dashdot")
 
     if args.temp > 0:
         hmin, hmax = axes[0].get_xlim()
@@ -146,7 +148,7 @@ def main():
 
     # Setting axes 0
     axes[0].set_title("Estimated energy", fontsize=fontsize)
-    axes[0].legend(frameon=False, fontsize=fontsize)  # Fat font, no frame
+    # axes[0].legend(frameon=False, fontsize=fontsize)  # Fat font, no frame
     axes[0].set_xlabel("r [A]", fontname=font, fontsize=fontsize)  # ylabel name, font is Hack, size 30pts
     axes[0].set_ylabel("E [kcal/mol]", fontname=font, fontsize=fontsize)  # ylabel name, font is Hack, size 30pts
     axes[0].tick_params(axis='both', which='major', labelsize=fontsize)  # Biggers ticks, bigger ticks label!
@@ -178,4 +180,3 @@ if __name__ == "__main__":
         main()
     except KeyboardInterrupt:
         raise SystemExit("User interruption.")
-