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
115
116
|
#!/usr/bin/env python3
# Copyright 2022 David Grote
#
# This file is part of WarpX.
#
# License: BSD-3-Clause-LBNL
"""
This script tests the quad lattice
The input file sets up a series of quadrupoles and propagates two particles through them.
One particle is in the X plane, the other the Y plane.
The final positions are compared to the analytic solutions.
The motion is slow enough that relativistic effects are ignored.
"""
import os
import sys
import numpy as np
from scipy.constants import c, e, m_e
import yt
yt.funcs.mylog.setLevel(0)
sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
import checksumAPI
filename = sys.argv[1]
ds = yt.load( filename )
ad = ds.all_data()
gamma_boost = float(ds.parameters.get('warpx.gamma_boost', 1.))
uz_boost = np.sqrt(gamma_boost*gamma_boost - 1.)*c
# Fetch the final particle position
xx_sim = ad['electron', 'particle_position_x'].v[0]
zz_sim = ad['electron', 'particle_position_z'].v[0]
ux_sim = ad['electron', 'particle_momentum_x'].v[0]/m_e
if gamma_boost > 1.:
# The simulation data is in the boosted frame.
# Transform the z position to the lab frame.
time = ds.current_time.value
zz_sim = gamma_boost*zz_sim + uz_boost*time;
# Fetch the quadrupole lattice data
quad_starts = []
quad_lengths = []
quad_strengths_E = []
z_location = 0.
def read_lattice(rootname, z_location):
lattice_elements = ds.parameters.get(f'{rootname}.elements').split()
for element in lattice_elements:
element_type = ds.parameters.get(f'{element}.type')
if element_type == 'drift':
length = float(ds.parameters.get(f'{element}.ds'))
z_location += length
elif element_type == 'quad':
length = float(ds.parameters.get(f'{element}.ds'))
quad_starts.append(z_location)
quad_lengths.append(length)
quad_strengths_E.append(float(ds.parameters.get(f'{element}.dEdx')))
z_location += length
elif element_type == 'lattice':
z_location = read_lattice(element, z_location)
return z_location
read_lattice('lattice', z_location)
# Fetch the initial position of the particle
x0 = [float(x) for x in ds.parameters.get('electron.single_particle_pos').split()]
ux0 = [float(x)*c for x in ds.parameters.get('electron.single_particle_vel').split()]
xx = x0[0]
zz = x0[2]
ux = ux0[0]
uz = ux0[2]
gamma = np.sqrt(uz**2/c**2 + 1.)
vz = uz/gamma
def applylens(x0, vx0, vz0, gamma, lens_length, lens_strength):
"""Use analytic solution of a particle with a transverse dependent field"""
kb0 = np.sqrt(e/(m_e*gamma*vz0**2)*abs(lens_strength))
if lens_strength >= 0.:
x1 = x0*np.cos(kb0*lens_length) + (vx0/vz0)/kb0*np.sin(kb0*lens_length)
vx1 = vz0*(-kb0*x0*np.sin(kb0*lens_length) + (vx0/vz0)*np.cos(kb0*lens_length))
else:
x1 = x0*np.cosh(kb0*lens_length) + (vx0/vz0)/kb0*np.sinh(kb0*lens_length)
vx1 = vz0*(+kb0*x0*np.sinh(kb0*lens_length) + (vx0/vz0)*np.cosh(kb0*lens_length))
return x1, vx1
# Integrate the particle using the analytic solution
for i in range(len(quad_starts)):
z_lens = quad_starts[i]
vx = ux/gamma
dt = (z_lens - zz)/vz
xx = xx + dt*vx
xx, vx = applylens(xx, vx, vz, gamma, quad_lengths[i], quad_strengths_E[i])
ux = gamma*vx
zz = z_lens + quad_lengths[i]
dt = (zz_sim - zz)/vz
vx = ux/gamma
xx = xx + dt*vx
# Compare the analytic to the simulated final values
print(f'Error in x position is {abs(np.abs((xx - xx_sim)/xx))}, which should be < 0.01')
print(f'Error in x velocity is {abs(np.abs((ux - ux_sim)/ux))}, which should be < 0.002')
assert abs(np.abs((xx - xx_sim)/xx)) < 0.01, Exception('error in x particle position')
assert abs(np.abs((ux - ux_sim)/ux)) < 0.002, Exception('error in x particle velocity')
test_name = os.path.split(os.getcwd())[1]
checksumAPI.evaluate_checksum(test_name, filename)
|
