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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
import os
import shutil
import numpy as np
import yt
import scipy.constants as scc
import glob
"""
This file is part of the suite of scripts to use LibEnsemble on top of WarpX
simulations. It reads the output plotfiles of a simulation and returns beam
quantity relevant for the LibEnsemble optimizations ('f') as well as other
beam quantities for convenience.
"""
yt.funcs.mylog.setLevel(50)
def slice_emittance(x, xp, g, b, w):
xpgb = xp * g * b
xav = np.average(x, weights = w)
xpav = np.average(xpgb, weights = w)
xstd2 = np.average((x - xav)**2, weights = w)
xpstd2 = np.average((xpgb - xpav)**2, weights = w)
xp2 = (np.average((x - xav) * (xpgb-xpav), weights = w))**2
em = np.sqrt(xstd2 * xpstd2 - xp2)
return em
def _beam_properties(filepath):
"""
Reads plotfile filepath and compute and return beam parameters
Parameters
----------
filepath : path to plotfile to read
"""
# Read beam quantities from plotfile
ds = yt.load(filepath)
if ('beam' in [i[0] for i in ds.field_list]):
ad = ds.all_data()
w = ad['beam', 'particle_weight'].v
if (w.shape[0] <= 200):
print('Insufficient particles in ',filepath)
return 0.0, 0.0, 0.0, 0.0
else:
x = ad['beam', 'particle_position_x'].v
z = ad['beam', 'particle_position_y'].v
ux = ad['beam', 'particle_momentum_x'].v/scc.m_e/scc.c
uy = ad['beam', 'particle_momentum_y'].v/scc.m_e/scc.c
uz = ad['beam', 'particle_momentum_z'].v/scc.m_e/scc.c
# Compute beam parameters
# Defined like that, the beam charge is > 0.
charge = np.sum(w) * scc.e
gamma = np.sqrt(1. + ux**2 + uy**2 + uz**2)
beta = np.sqrt(1.0 - 1.0 / gamma**2)
energy_MeV = scc.physical_constants["electron mass energy equivalent in MeV"][0] * (gamma - 1.)
energy_avg = np.average(energy_MeV, weights = w)
energy_std = np.sqrt(np.average((energy_MeV - energy_avg)**2, weights = w)) / energy_avg * 100
nslices = 20
zslices = np.linspace(np.min(z), np.max(z), nslices+1)
exlist = np.zeros(nslices)
for slicei in range(nslices):
cond = [(z > zslices[slicei]) & (z < zslices[slicei+1])][0]
wslice = w[cond]
if (wslice.shape[0] > 10):
xslice = x[cond]
xpslice = ux[cond]/uz[cond]
gslice = gamma[cond]
bslice = beta[cond]
exlist[slicei] = slice_emittance(xslice, xpslice, gslice, bslice, wslice)
emittance = np.mean(exlist)
return charge, energy_avg, energy_std, emittance
else:
print('No beam particles in ',filepath)
return 0.0, 0.0, 0.0, 0.0
def read_sim_output(workdir):
"""
Return optimizing quantity 'f' and other parameters for convenience.
Parameters
----------
workdir : Path to directory where the simulation ran.
"""
# Get beam properties at the beginning of the run
file_list = glob.glob('diags/plotfiles/plt?????')
if (len(file_list) <2):
print(workdir,' did not have final plotfile')
return np.array([np.nan, np.nan, np.nan, np.nan, np.nan])
file_list.sort()
datafile = file_list[0]
filepath = os.path.join(workdir, datafile)
charge_i, _, _, emittance_i = _beam_properties(filepath)
# Get beam properties at the end of the run
datafile = file_list[-1]
filepath = os.path.join(workdir, datafile)
charge_f, energy_avg, energy_std, emittance_f = _beam_properties(filepath)
if (charge_f > 0.0):
# delete simulation results, just to have smaller data
shutil.rmtree('diags')
# Build a quantity to minimize (f) that encompasses
# emittance AND charge loss 1% charge loss has the
# same impact as doubling the initial emittance.
# we minimize f!
f = emittance_f + emittance_i*(1.-charge_f/charge_i)*100
warpx_out = np.array([f, energy_std, energy_avg, charge_f, emittance_f])
return warpx_out
else:
return np.array([np.nan, np.nan, np.nan, np.nan, np.nan])
|
