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
|
#!/usr/bin/env python3
# Copyright 2019-2020 Luca Fedeli, Maxence Thevenet, Revathi Jambunathan
#
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL
'''
Analysis script of a WarpX simulation of rigid injection in a boosted frame.
A Gaussian electron beam starts from -5 microns, propagates rigidly up to
20 microns after which it expands due to emittance only (the focal position is
20 microns). The beam width is measured after ~50 microns, and compared with
the theory (with a 1% relative error allowed).
The simulation runs in a boosted frame, and the analysis is done in the lab
frame, i.e., on the back-transformed diagnostics.
'''
import os
import sys
import numpy as np
import openpmd_api as io
import read_raw_data
from scipy.constants import m_e
import yt
yt.funcs.mylog.setLevel(0)
sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
import checksumAPI
filename = sys.argv[1]
# Tolerances to check consistency between legacy BTD and new BTD
rtol = 1e-16
atol = 1e-16
# Read data from legacy back-transformed diagnostics
snapshot = './lab_frame_data/snapshots/snapshot00001'
x_legacy = read_raw_data.get_particle_field(snapshot, 'beam', 'x')
z_legacy = read_raw_data.get_particle_field(snapshot, 'beam', 'z')
ux_legacy = read_raw_data.get_particle_field(snapshot, 'beam', 'ux')
uy_legacy = read_raw_data.get_particle_field(snapshot, 'beam', 'uy')
uz_legacy = read_raw_data.get_particle_field(snapshot, 'beam', 'uz')
# Read data from new back-transformed diagnostics (plotfile)
ds_plotfile = yt.load(filename)
x_plotfile = ds_plotfile.all_data()['beam', 'particle_position_x'].v
z_plotfile = ds_plotfile.all_data()['beam', 'particle_position_y'].v
ux_plotfile = ds_plotfile.all_data()['beam', 'particle_momentum_x'].v
uy_plotfile = ds_plotfile.all_data()['beam', 'particle_momentum_y'].v
uz_plotfile = ds_plotfile.all_data()['beam', 'particle_momentum_z'].v
# Read data from new back-transformed diagnostics (openPMD)
series = io.Series("./diags/diag2/openpmd_%T.h5", io.Access.read_only)
ds_openpmd = series.iterations[1]
x_openpmd = ds_openpmd.particles['beam']['position']['x'][:]
z_openpmd = ds_openpmd.particles['beam']['position']['z'][:]
ux_openpmd = ds_openpmd.particles['beam']['momentum']['x'][:]
uy_openpmd = ds_openpmd.particles['beam']['momentum']['y'][:]
uz_openpmd = ds_openpmd.particles['beam']['momentum']['z'][:]
series.flush()
# Sort and compare arrays to check consistency between legacy BTD and new BTD (plotfile)
assert(np.allclose(np.sort(x_legacy), np.sort(x_plotfile), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(z_legacy), np.sort(z_plotfile), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(ux_legacy*m_e), np.sort(ux_plotfile), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(uy_legacy*m_e), np.sort(uy_plotfile), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(uz_legacy*m_e), np.sort(uz_plotfile), rtol=rtol, atol=atol))
# Sort and compare arrays to check consistency between legacy BTD and new BTD (openPMD)
assert(np.allclose(np.sort(x_legacy), np.sort(x_openpmd), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(z_legacy), np.sort(z_openpmd), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(ux_legacy*m_e), np.sort(ux_openpmd), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(uy_legacy*m_e), np.sort(uy_openpmd), rtol=rtol, atol=atol))
assert(np.allclose(np.sort(uz_legacy*m_e), np.sort(uz_openpmd), rtol=rtol, atol=atol))
# Initial parameters
z0 = 20.e-6
x0 = 1.e-6
theta0 = np.arcsin(0.1)
# Theoretical beam width after propagation with rigid injection
z = np.mean(z_legacy)
x = np.std(x_legacy)
print(f'Beam position = {z}')
print(f'Beam width = {x}')
xth = np.sqrt(x0**2 + (z-z0)**2 * theta0**2)
err = np.abs((x-xth) / xth)
tol = 1e-2
print(f'error = {err}')
print(f'tolerance = {tol}')
assert(err < tol)
test_name = os.path.split(os.getcwd())[1]
checksumAPI.evaluate_checksum(test_name, filename)
|
