1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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()
|
