1 files changed, 83 insertions, 0 deletions
diff --git a/Tools/parallel_postproc.py b/Tools/parallel_postproc.py
new file mode 100644
index 000000000..0942118ac
--- /dev/null
+++ b/Tools/parallel_postproc.py
@@ -0,0 +1,83 @@
+import matplotlib
+matplotlib.use('Agg')
+from mpi4py import MPI
+import glob, read_raw_data
+import numpy as np
+import scipy.constants as scc
+import  matplotlib.pyplot as plt
+
+species = 'electrons'
+fieldname = 'Ey'
+lambda0 = .81e-6
+
+# --- custom functions --- #
+# ------------------------ #
+omega0 = 2*np.pi*scc.c/lambda0
+# Read field fieldname and return normalized max
+def get_a0(res_dir, snapshot):
+    header = res_dir + '/Header'
+    print( snapshot )
+    allrd, info = read_raw_data.read_lab_snapshot(snapshot, header)
+    F = allrd[ fieldname ]
+    return info['z'][-1], np.max(np.abs(F)) * scc.e/(scc.m_e*omega0*scc.c)
+# Convert elements of a list to numpy arrays
+def convert_to_np_array(list_in):
+    list_out = []
+    for elem in list_in:
+        list_out.append( np.array( elem ) )
+    return list_out
+
+# --- MPI parallelization --- #
+# --------------------------- #
+# Get ordered list of snapshot files
+res_dir = './lab_frame_data/';
+file_list = glob.glob(res_dir + '/snapshot?????')
+file_list.sort()
+# Total number of files
+nfiles = len(file_list)
+# Each file has a unique number
+number_list = range(nfiles)
+comm_world = MPI.COMM_WORLD
+me = comm_world.Get_rank()
+nrank = comm_world.Get_size()
+# List of files to process for current proc
+my_list = file_list[ (me*nfiles)/nrank : ((me+1)*nfiles)/nrank ]
+# List if file numbers for current proc
+my_number_list = number_list[ (me*nfiles)/nrank : ((me+1)*nfiles)/nrank ]
+
+# --- Run parallel analysis --- #
+# ----------------------------- #
+# Each MPI rank reads roughly (nb snapshots)/(nb ranks) snapshots.
+# Works with any number of snapshots.
+for count, filename in enumerate(my_list):
+    zwin, a0 = get_a0( res_dir, filename )
+    uzlab = read_raw_data.get_particle_field(filename, species, 'uz')/scc.c
+    select_particles = (uzlab > 5.)
+    uzlab = uzlab[ select_particles ]
+    uzmean = np.mean(uzlab)
+
+# --- gather and rank 0 plots --- #
+# ------------------------------- #
+# Gather particle quantities to rank 0 to plot history of quantities.
+UZMEAN = comm_world.gather(uzmean, root=0)
+ZWIN = comm_world.gather(zwin, root=0)
+A0 = comm_world.gather(a0, root=0)
+# Rank 0 does the plot.
+if me == 0:
+    # Convert to numpy arrays
+    UZMEAN, ZWIN, A0 = convert_to_np_array([UZMEAN, ZWIN, A0])
+    # Plot and save
+    fig = plt.figure()
+    plt.subplot(2,1,1)
+    plt.plot(ZWIN*1.e3, UZMEAN)
+    plt.xlabel('z (mm)')
+    plt.ylabel('uz')
+    plt.title( 'beam energy' )
+    plt.grid()
+    plt.subplot(2,1,2)
+    plt.plot(ZWIN*1.e3, A0)
+    plt.ylabel('a0')
+    plt.title( 'beam propag angle' )
+    plt.grid()
+    fig.savefig('./image.png', bbox_inches='tight')
+    plt.close()