From b876b47dc9dd9ed8c940668db0dd387a24594280 Mon Sep 17 00:00:00 2001
From: Dave Grote <grote1@llnl.gov>
Date: Wed, 18 Sep 2019 13:33:12 -0700
Subject: Various clean up in plot_parallel.py

---
 Tools/plot_parallel.py | 61 ++++++++++++++++++++++++++------------------------
 1 file changed, 32 insertions(+), 29 deletions(-)

(limited to 'Tools/plot_parallel.py')

diff --git a/Tools/plot_parallel.py b/Tools/plot_parallel.py
index 56b243aca..d9efaa93a 100644
--- a/Tools/plot_parallel.py
+++ b/Tools/plot_parallel.py
@@ -1,8 +1,12 @@
-import os, glob, matplotlib, sys, argparse
-import yt ; yt.funcs.mylog.setLevel(50)
+import os
+import glob
+import matplotlib
+import sys
+import argparse
+import yt
+yt.funcs.mylog.setLevel(50)
 import numpy as np
 import matplotlib.pyplot as plt
-import scipy.constants as scc
 
 '''
 This script loops over all WarpX plotfiles in a directory and, for each
@@ -28,22 +32,20 @@ To get help, run
 
 # Parse command line for options.
 parser = argparse.ArgumentParser()
-parser.add_argument('--path', dest='path', default='.',
+parser.add_argument('--path', default='diags/plotfiles',
                     help='path to plotfiles. Plotfiles names must be plt?????')
-parser.add_argument('--plotlib', dest='plotlib', default='yt',
+parser.add_argument('--plotlib', default='yt',
                     choices=['yt','matplotlib'],
                     help='Plotting library to use')
-parser.add_argument('--field', dest='field', default='Ez',
+parser.add_argument('--field', default='Ez',
                     help='Which field to plot, e.g., Ez, By, jx or rho. The central slice in y is plotted')
-parser.add_argument('--pjump', dest='pjump', default=20,
+parser.add_argument('--pjump', default=20,
                     help='When plotlib=matplotlib, we plot every pjump particle')
-parser.add_argument('--use_vmax', dest='use_vmax', default=False,
-                    help='Whether to put bounds to field colormap')
-parser.add_argument('--vmax', dest='vmax', default=1.e12,
-                    help='If use_vmax=True, the colormab will have bounds [-vamx, vmax]')
-parser.add_argument('--slicewidth', dest='slicewidth', default=10.e-6,
+parser.add_argument('--vmax', type=float, default=None,
+                    help='If specified, the colormap will have bounds [-vmax, vmax]')
+parser.add_argument('--slicewidth', default=10.e-6,
                     help='Only particles with -slicewidth/2<y<slicewidth/2 are plotted')
-parser.add_argument('--parallel', dest='parallel', action='store_true', default=False,
+parser.add_argument('--parallel', action='store_true', default=False,
                     help='whether or not to do the analysis in parallel (e.g., 1 plotfile per MPI rank)')
 parser.add_argument('--species', dest='pslist', nargs='+', type=str, default=None,
                     help='Species to be plotted, e.g., " --species beam plasma_e ". By default, all species in the simulation are shown')
@@ -51,7 +53,10 @@ parser.add_argument('--plot_max_evolution', type=str, default=None,
                     help='Quantity to plot the max of across all data files')
 args = parser.parse_args()
 
+plotlib = args.plotlib
 plot_max_evolution = args.plot_max_evolution
+vmax = args.vmax
+
 # Sanity check
 if int(sys.version[0]) != 3:
     print('WARNING: Parallel analysis was only tested with Python3')
@@ -68,14 +73,14 @@ def get_species(a_file_list):
     if args.pslist is not None:
         return args.pslist
     # otherwise, loop over all plotfiles to get particle species list
-    pslist = []
+    psset = set()
     for filename in a_file_list:
         ds = yt.load( filename )
-        # get list of species in current plotfile
-        pslist_plotfile = list( set( [x[0] for x in ds.field_list
-                                      if x[1][:9]=='particle_'] ) )
-        # append species in current plotfile to pslist, and uniquify
-        pslist = list( set( pslist + pslist_plotfile ) )
+        for ps in ds.particle_types:
+            if ps == 'all':
+                continue
+            psset.add(ps)
+    pslist = list(psset)
     pslist.sort()
     return pslist
 
@@ -101,11 +106,13 @@ def plot_snapshot(filename):
         plt.colorbar()
         plt.xlim(ds.domain_left_edge[dim-1], ds.domain_right_edge[dim-1])
         plt.ylim(ds.domain_left_edge[0], ds.domain_right_edge[0])
-        if args.use_vmax:
-            plt.clim(-args.vmax, args.vmax)
+        if vmax is not None:
+            plt.clim(-vmax, vmax)
     if plotlib == 'yt':
         # Directly plot with yt
         sl = yt.SlicePlot(ds, yt_slicedir[dim], args.field, aspect=yt_aspect[dim])
+        if vmax is not None:
+            sl.set_zlim(-vmax, vmax)
 
     # Plot particle quantities
     for ispecies, pspecies in enumerate(pslist):
@@ -136,15 +143,12 @@ def plot_snapshot(filename):
     if plotlib == 'matplotlib':
         plt.xlabel('z (m)')
         plt.ylabel('x (m)')
-        plt.title(args.field + ' at iteration ' + str(iteration) +
-                  ', time = ' + str(ds.current_time))
-        plt.savefig(args.path + '/plt_' + args.field + '_' + plotlib + '_' +
-                    str(iteration).zfill(5) + '.png', bbox_inches='tight', dpi=300)
+        plt.title('%s at iteration %d, time = %e s'%(args.field, iteration, ds.current_time))
+        plt.savefig('plt_%s_%s_%05d.png'%(args.field, plotlib, iteration), bbox_inches='tight', dpi=300)
         plt.close()
     if plotlib == 'yt':
         sl.annotate_grids()
-        sl.save(args.path + '/plt_' + args.field + '_' + plotlib + '_' +
-                str(iteration).zfill(5) + '.png')
+        sl.save('plt_%s_%s_%05d.png'%(args.field, plotlib, iteration))
 
 # Compute max of field a_field in plotfile filename
 def get_field_max( filename, a_field ):
@@ -170,8 +174,7 @@ def plot_field_max(zwin_arr, maxF_arr):
 ### Analysis ###
 
 # Get list of plotfiles
-plotlib  = args.plotlib
-file_list = glob.glob(args.path + '/plt?????')
+file_list = glob.glob(os.path.join(args.path, 'plt?????'))
 file_list.sort()
 nfiles = len(file_list)
 
-- 
cgit v1.2.3